#!/bin/bash
################################################################################
# Script Name: install-snort.sh
# Description: Automated Snort 3 IDS/IPS installation from source with
#              rule management on Ubuntu/Debian and RHEL/Rocky/Alma/Fedora
#
# Author:  Phil Connor
# Contact: contact@mylinux.work
# Website: https://mylinux.work
# License: MIT
# Version: 1.01
#
# Usage:
#   sudo ./install-snort.sh
#   sudo ./install-snort.sh --iface eth0 --oinkcode YOUR_CODE
#   sudo ./install-snort.sh --community-rules
#   sudo ./install-snort.sh --dry-run
#
################################################################################

set -euo pipefail

# ============================================================================
# DEFAULTS
# ============================================================================

SNORT_VERSION="3.1.84.0"
DAQ_VERSION="3.0.14"
IFACE=""
OINKCODE=""
COMMUNITY_RULES=true
REGISTERED_RULES=false
HOME_NET="[192.168.0.0/16,10.0.0.0/8,172.16.0.0/12]"
INSTALL_DIR="/usr/local"
CONFIG_DIR="/etc/snort"
LOG_DIR="/var/log/snort"
RULE_DIR="/etc/snort/rules"
DRY_RUN=false
UNINSTALL=false
SKIP_BUILD=false

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
NC='\033[0m'

# ============================================================================
# HELPER FUNCTIONS
# ============================================================================

log_info()  { echo -e "${GREEN}[INFO]${NC} $*"; }
log_warn()  { echo -e "${YELLOW}[WARN]${NC} $*"; }
log_error() { echo -e "${RED}[ERROR]${NC} $*" >&2; }
log_step()  { echo -e "${CYAN}[STEP]${NC} $*"; }

show_usage() {
    cat <<EOF
Usage: $0 [OPTIONS]

Installs Snort 3 IDS/IPS from source with rule management.

OPTIONS:
    --iface IFACE           Network interface to monitor (auto-detected if omitted)
    --snort-version VER     Snort version to install (default: $SNORT_VERSION)
    --daq-version VER       DAQ version (default: $DAQ_VERSION)
    --oinkcode CODE         Snort.org Oinkcode for registered/subscriber rules
    --community-rules       Install community rules (default, free)
    --registered-rules      Download registered rules (requires --oinkcode)
    --home-net CIDR         HOME_NET definition (default: RFC1918 ranges)
    --install-dir DIR       Installation prefix (default: /usr/local)
    --skip-build            Skip build, configure only (re-run after initial install)
    --uninstall             Remove Snort and configuration
    --dry-run               Show what would be done without executing
    -h, --help              Show this help message

EXAMPLES:
    $0                                          # Community rules, auto-detect interface
    $0 --iface eth0                             # Specific interface
    $0 --oinkcode ABC123 --registered-rules     # Registered rules with Oinkcode
    $0 --skip-build                             # Reconfigure existing install
    $0 --dry-run                                # Preview changes

EOF
    exit 0
}

parse_args() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h|--help)              show_usage ;;
            --iface)                IFACE="$2"; shift 2 ;;
            --snort-version)        SNORT_VERSION="$2"; shift 2 ;;
            --daq-version)          DAQ_VERSION="$2"; shift 2 ;;
            --oinkcode)             OINKCODE="$2"; shift 2 ;;
            --community-rules)      COMMUNITY_RULES=true; shift ;;
            --registered-rules)     REGISTERED_RULES=true; shift ;;
            --home-net)             HOME_NET="$2"; shift 2 ;;
            --install-dir)          INSTALL_DIR="$2"; shift 2 ;;
            --skip-build)           SKIP_BUILD=true; shift ;;
            --uninstall)            UNINSTALL=true; shift ;;
            --dry-run)              DRY_RUN=true; shift ;;
            *) log_error "Unknown option: $1"; exit 1 ;;
        esac
    done
}

check_root() {
    if [[ $EUID -ne 0 ]]; then
        log_error "This script must be run as root (sudo)"
        exit 1
    fi
}

detect_os() {
    if [[ -f /etc/os-release ]]; then
        . /etc/os-release
        OS_ID="$ID"
        OS_VERSION="${VERSION_ID%%.*}"
        OS_FAMILY=""
        case "$OS_ID" in
            ubuntu|debian)          OS_FAMILY="debian" ;;
            rhel|centos|rocky|almalinux|fedora) OS_FAMILY="rhel" ;;
            *) log_error "Unsupported OS: $OS_ID"; exit 1 ;;
        esac
        log_info "Detected OS: $PRETTY_NAME ($OS_FAMILY)"
    else
        log_error "Cannot detect OS"
        exit 1
    fi
}

detect_interface() {
    if [[ -z "$IFACE" ]]; then
        IFACE=$(ip route show default 2>/dev/null | awk '/default/ {print $5; exit}')
        if [[ -z "$IFACE" ]]; then
            IFACE=$(ip -o link show up 2>/dev/null | awk -F': ' '!/lo/{print $2; exit}')
        fi
        if [[ -z "$IFACE" ]]; then
            log_error "Cannot auto-detect interface. Use --iface"
            exit 1
        fi
        log_info "Auto-detected interface: $IFACE"
    fi
}

# ============================================================================
# BUILD DEPENDENCIES
# ============================================================================

install_deps_debian() {
    log_step "Installing build dependencies (Debian/Ubuntu)..."
    apt-get update -qq
    apt-get install -y -qq \
        build-essential cmake pkg-config \
        libhwloc-dev libluajit-5.1-dev libssl-dev \
        libpcap-dev libpcre3-dev zlib1g-dev \
        libdumbnet-dev liblzma-dev libsafec-dev \
        libunwind-dev uuid-dev \
        flex bison \
        libflatbuffers-dev flatbuffers-compiler \
        libhyperscan-dev \
        libjemalloc-dev \
        git wget curl jq \
        cpputest libcpputest-dev
}

install_deps_rhel() {
    log_step "Installing build dependencies (RHEL/Rocky/Alma)..."
    if [[ "$OS_ID" != "fedora" ]]; then
        dnf install -y -q epel-release
        dnf config-manager --set-enabled crb 2>/dev/null || \
            dnf config-manager --set-enabled powertools 2>/dev/null || true
    fi
    dnf groupinstall -y -q "Development Tools"
    dnf install -y -q \
        cmake3 pkgconfig \
        hwloc-devel luajit-devel openssl-devel \
        libpcap-devel pcre-devel zlib-devel \
        libdnet-devel xz-devel libsafec-devel \
        libunwind-devel uuid-devel \
        flex bison \
        flatbuffers-devel flatbuffers-compiler \
        hyperscan-devel \
        jemalloc-devel \
        git wget curl jq
}

install_dependencies() {
    case "$OS_FAMILY" in
        debian) install_deps_debian ;;
        rhel)   install_deps_rhel ;;
    esac
}

# ============================================================================
# BUILD & INSTALL
# ============================================================================

build_libdaq() {
    log_step "Building libdaq $DAQ_VERSION..."

    local build_dir="/tmp/snort-build"
    mkdir -p "$build_dir"
    cd "$build_dir"

    if [[ ! -d "libdaq-$DAQ_VERSION" ]]; then
        wget -q "https://github.com/snort3/libdaq/archive/refs/tags/v$DAQ_VERSION.tar.gz" \
            -O "libdaq-$DAQ_VERSION.tar.gz"
        tar xzf "libdaq-$DAQ_VERSION.tar.gz"
    fi

    cd "libdaq-$DAQ_VERSION"
    ./bootstrap 2>/dev/null || true
    ./configure --prefix="$INSTALL_DIR"
    make -j"$(nproc)"
    make install

    ldconfig
    log_info "libdaq $DAQ_VERSION installed"
}

build_snort() {
    log_step "Building Snort $SNORT_VERSION..."

    local build_dir="/tmp/snort-build"
    mkdir -p "$build_dir"
    cd "$build_dir"

    if [[ ! -d "snort3-$SNORT_VERSION" ]]; then
        wget -q "https://github.com/snort3/snort3/archive/refs/tags/$SNORT_VERSION.tar.gz" \
            -O "snort3-$SNORT_VERSION.tar.gz"
        tar xzf "snort3-$SNORT_VERSION.tar.gz"
    fi

    cd "snort3-$SNORT_VERSION"
    mkdir -p build && cd build

    cmake .. \
        -DCMAKE_INSTALL_PREFIX="$INSTALL_DIR" \
        -DENABLE_LARGE_PCAP=ON \
        -DENABLE_JEMALLOC=ON \
        2>&1 | tail -5

    make -j"$(nproc)"
    make install

    ldconfig
    log_info "Snort $SNORT_VERSION installed to $INSTALL_DIR"

    # Verify
    "$INSTALL_DIR/bin/snort" -V 2>&1 | head -3
}

# ============================================================================
# CONFIGURATION
# ============================================================================

create_snort_user() {
    if ! id snort &>/dev/null; then
        groupadd -r snort 2>/dev/null || true
        useradd -r -g snort -s /sbin/nologin -d /var/log/snort snort 2>/dev/null || true
        log_info "Created snort user and group"
    fi
}

configure_directories() {
    log_step "Creating directory structure..."

    mkdir -p "$CONFIG_DIR"
    mkdir -p "$CONFIG_DIR/rules"
    mkdir -p "$CONFIG_DIR/builtin_rules"
    mkdir -p "$CONFIG_DIR/so_rules"
    mkdir -p "$CONFIG_DIR/lists"
    mkdir -p "$LOG_DIR"
    mkdir -p /var/lib/snort

    chown -R snort:snort "$LOG_DIR"
    chown -R snort:snort /var/lib/snort
}

configure_snort() {
    log_step "Configuring Snort..."

    # Copy default config if not present
    if [[ ! -f "$CONFIG_DIR/snort.lua" ]]; then
        if [[ -f "$INSTALL_DIR/etc/snort/snort.lua" ]]; then
            cp "$INSTALL_DIR/etc/snort/snort.lua" "$CONFIG_DIR/snort.lua"
            cp "$INSTALL_DIR/etc/snort/snort_defaults.lua" "$CONFIG_DIR/" 2>/dev/null || true
            cp "$INSTALL_DIR/etc/snort/file_magic.lua" "$CONFIG_DIR/" 2>/dev/null || true
        fi
    fi

    # Create local.lua overrides
    cat > "$CONFIG_DIR/local.lua" <<LUAEOF
-- Local Snort configuration overrides
-- Generated by install-snort.sh

HOME_NET = '$HOME_NET'
EXTERNAL_NET = '!\$HOME_NET'

-- IDS mode configuration
ips = {
    enable_builtin_rules = true,
    include = RULE_PATH .. '/snort3-community.rules',
    variables = default_variables,
}

-- EVE-style JSON alert output
alert_json = {
    file = true,
    limit = 100,
    fields = 'timestamp msg src_addr src_port dst_addr dst_port proto action',
}

-- Unified2 output (for Barnyard2 or legacy tools)
-- alert_unified2 = { limit = 128 }

-- Logging
output = {
    logdir = '$LOG_DIR',
}
LUAEOF

    log_info "Snort configuration written to $CONFIG_DIR/"
}

# ============================================================================
# RULES
# ============================================================================

download_rules() {
    log_step "Downloading Snort rules..."

    if [[ "$COMMUNITY_RULES" == true ]]; then
        log_info "Downloading community rules..."
        local community_url="https://www.snort.org/downloads/community/snort3-community-rules.tar.gz"
        wget -q "$community_url" -O /tmp/snort3-community-rules.tar.gz || {
            log_warn "Failed to download community rules — trying alternative URL"
            wget -q "https://www.snort.org/downloads/community/community-rules.tar.gz" \
                -O /tmp/snort3-community-rules.tar.gz || {
                log_error "Cannot download community rules"
                return 1
            }
        }
        tar xzf /tmp/snort3-community-rules.tar.gz -C /tmp/
        find /tmp/ -name '*.rules' -path '*community*' -exec cp {} "$RULE_DIR/" \;
        log_info "Community rules installed to $RULE_DIR/"
    fi

    if [[ "$REGISTERED_RULES" == true ]]; then
        if [[ -z "$OINKCODE" ]]; then
            log_warn "Registered rules require --oinkcode — skipping"
        else
            log_info "Downloading registered rules with Oinkcode..."
            local reg_url="https://www.snort.org/reg-rules/snortrules-snapshot-31840.tar.gz/$OINKCODE"
            wget -q "$reg_url" -O /tmp/snort-registered-rules.tar.gz || {
                log_error "Failed to download registered rules (check Oinkcode)"
                return 1
            }
            tar xzf /tmp/snort-registered-rules.tar.gz -C "$RULE_DIR/"
            log_info "Registered rules installed to $RULE_DIR/"
        fi
    fi

    # Count rules
    local rule_count
    rule_count=$(grep -rch '^\(alert\|drop\|reject\|pass\)' "$RULE_DIR/" 2>/dev/null || true | awk '{s+=$1} END {print s+0}')
    log_info "Total rules installed: $rule_count"
}

setup_rule_update_cron() {
    log_step "Setting up weekly rule update cron..."

    cat > /etc/cron.d/snort-rule-update <<CRONEOF
# Update Snort rules weekly (Sunday 03:00) and restart
0 3 * * 0 root /usr/local/bin/install-snort.sh --skip-build --community-rules 2>/dev/null && systemctl restart snort 2>/dev/null
CRONEOF

    chmod 644 /etc/cron.d/snort-rule-update
    log_info "Weekly rule update cron job created (Sunday 03:00)"
}

# ============================================================================
# SYSTEMD SERVICE
# ============================================================================

create_service() {
    log_step "Creating Snort systemd service..."

    cat > /etc/systemd/system/snort.service <<SVCEOF
[Unit]
Description=Snort 3 IDS/IPS
After=network-online.target
Wants=network-online.target

[Service]
Type=simple
ExecStart=$INSTALL_DIR/bin/snort -c $CONFIG_DIR/snort.lua --plugin-path $INSTALL_DIR/lib/snort/daqs -i $IFACE -l $LOG_DIR -D -u snort -g snort
ExecReload=/bin/kill -SIGHUP \$MAINPID
Restart=on-failure
RestartSec=10
LimitNOFILE=65536

[Install]
WantedBy=multi-user.target
SVCEOF

    systemctl daemon-reload
    systemctl enable snort
    systemctl restart snort

    sleep 3

    if systemctl is-active --quiet snort; then
        log_info "Snort is running"
    else
        log_warn "Snort failed to start — check: journalctl -u snort"
        systemctl status snort --no-pager 2>/dev/null || true
    fi
}

# ============================================================================
# VALIDATION
# ============================================================================

validate_installation() {
    log_step "Validating Snort installation..."

    "$INSTALL_DIR/bin/snort" -c "$CONFIG_DIR/snort.lua" --warn-all -T 2>&1 | tail -5 || {
        log_warn "Config validation produced warnings"
    }

    echo ""
    log_info "===== Installation Summary ====="
    log_info "Snort version:   $SNORT_VERSION"
    log_info "DAQ version:     $DAQ_VERSION"
    log_info "Interface:       $IFACE"
    log_info "Config:          $CONFIG_DIR/snort.lua"
    log_info "Rules:           $RULE_DIR/"
    log_info "Logs:            $LOG_DIR/"
    log_info "HOME_NET:        $HOME_NET"
    echo ""
}

# ============================================================================
# UNINSTALL
# ============================================================================

uninstall_snort() {
    log_step "Uninstalling Snort..."

    systemctl stop snort 2>/dev/null || true
    systemctl disable snort 2>/dev/null || true
    rm -f /etc/systemd/system/snort.service
    rm -f /etc/cron.d/snort-rule-update
    systemctl daemon-reload

    rm -rf "$INSTALL_DIR/bin/snort"
    rm -rf "$INSTALL_DIR/lib/snort"
    rm -rf "$INSTALL_DIR/etc/snort"

    log_info "Snort binaries removed"
    log_info "Config ($CONFIG_DIR/) and logs ($LOG_DIR/) left intact"
    log_info "Remove manually if no longer needed"
}

# ============================================================================
# DRY RUN
# ============================================================================

dry_run() {
    echo ""
    log_info "===== DRY RUN — No changes will be made ====="
    echo ""
    log_info "OS:               $PRETTY_NAME"
    log_info "Snort version:    $SNORT_VERSION"
    log_info "DAQ version:      $DAQ_VERSION"
    log_info "Interface:        $IFACE"
    log_info "HOME_NET:         $HOME_NET"
    log_info "Install prefix:   $INSTALL_DIR"
    log_info "Config dir:       $CONFIG_DIR"
    log_info "Log dir:          $LOG_DIR"
    log_info "Community rules:  $COMMUNITY_RULES"
    log_info "Registered rules: $REGISTERED_RULES"
    echo ""
    log_info "Actions that would be performed:"
    echo "  1. Install build dependencies"
    echo "  2. Build libdaq $DAQ_VERSION from source"
    echo "  3. Build Snort $SNORT_VERSION from source"
    echo "  4. Create snort user and group"
    echo "  5. Create directory structure"
    echo "  6. Write Snort configuration"
    echo "  7. Download and install rules"
    echo "  8. Create weekly rule update cron job"
    echo "  9. Create and start systemd service"
    echo ""
}

# ============================================================================
# MAIN
# ============================================================================

main() {
    parse_args "$@"
    check_root
    detect_os
    detect_interface

    if [[ "$UNINSTALL" == true ]]; then
        uninstall_snort
        exit 0
    fi

    if [[ "$DRY_RUN" == true ]]; then
        dry_run
        exit 0
    fi

    echo ""
    log_info "===== Snort 3 IDS/IPS Installer ====="
    echo ""

    if [[ "$SKIP_BUILD" == false ]]; then
        install_dependencies
        build_libdaq
        build_snort
    fi

    create_snort_user
    configure_directories
    configure_snort
    download_rules
    setup_rule_update_cron
    create_service
    validate_installation

    echo ""
    log_info "===== Installation Complete ====="
    log_info "View alerts:     tail -f $LOG_DIR/alert_json.txt | jq ."
    log_info "Test config:     snort -c $CONFIG_DIR/snort.lua -T"
    log_info "Service status:  systemctl status snort"
    echo ""
}

main "$@"