#!/bin/bash
################################################################################
# Script Name: install-suricata.sh
# Description: Automated Suricata IDS/IPS installation with multi-source
#              rule selection 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-suricata.sh
#   sudo ./install-suricata.sh --iface eth0 --rules et/open,oisf/trafficid
#   sudo ./install-suricata.sh --ips --iface ens192
#   sudo ./install-suricata.sh --dry-run
#
################################################################################

set -euo pipefail

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

IFACE=""
MODE="ids"
RULE_SOURCES="et/open"
HOME_NET="[192.168.0.0/16,10.0.0.0/8,172.16.0.0/12]"
THREADS="auto"
EVE_LOG="/var/log/suricata/eve.json"
SURICATA_CONF="/etc/suricata/suricata.yaml"
DRY_RUN=false
UNINSTALL=false
SKIP_RULES=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 and configures Suricata IDS/IPS with multi-source rule selection.

OPTIONS:
    --iface IFACE       Network interface to monitor (auto-detected if omitted)
    --ids               Run in IDS mode (default — passive monitoring)
    --ips               Run in IPS mode (inline, drops malicious traffic)
    --rules SOURCES     Comma-separated rule sources (default: et/open)
    --home-net CIDR     HOME_NET definition (default: RFC1918 ranges)
    --threads N         Worker threads (default: auto-detect from CPU cores)
    --eve-log PATH      EVE JSON log path (default: /var/log/suricata/eve.json)
    --skip-rules        Skip rule download (install only)
    --uninstall         Remove Suricata and configuration
    --dry-run           Show what would be done without executing
    -h, --help          Show this help message

RULE SOURCES (use with --rules, comma-separated):
    et/open             Emerging Threats Open (default, free)
    et/pro CODE         Emerging Threats Pro (requires Oinkcode)
    oisf/trafficid      OISF Traffic ID rules
    ptresearch/attackdetection  Positive Technologies attack detection
    sslbl/ssl-fp-blacklist      Abuse.ch SSL fingerprint blacklist
    sslbl/ja3-fingerprints      Abuse.ch JA3 fingerprint blacklist
    etnetera/aggressive         Etnetera aggressive IP blacklist
    tgreen/hunting              tgreen hunting rules
    malsilo/win-malware         Malsilo Windows malware rules
    stamus/lateral              Stamus lateral movement detection

EXAMPLES:
    $0                                           # Auto-detect, ET Open, IDS mode
    $0 --iface eth0                              # Specific interface
    $0 --rules et/open,oisf/trafficid            # Multiple rule sources
    $0 --rules et/open,sslbl/ssl-fp-blacklist    # ET Open + SSL blacklist
    $0 --ips --iface ens192                      # IPS inline mode
    $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 ;;
            --ids)          MODE="ids"; shift ;;
            --ips)          MODE="ips"; shift ;;
            --rules)        RULE_SOURCES="$2"; shift 2 ;;
            --home-net)     HOME_NET="$2"; shift 2 ;;
            --threads)      THREADS="$2"; shift 2 ;;
            --eve-log)      EVE_LOG="$2"; shift 2 ;;
            --skip-rules)   SKIP_RULES=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 (missing /etc/os-release)"
        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 network interface. Use --iface"
            exit 1
        fi
        log_info "Auto-detected interface: $IFACE"
    else
        if ! ip link show "$IFACE" &>/dev/null; then
            log_error "Interface $IFACE does not exist"
            exit 1
        fi
        log_info "Using interface: $IFACE"
    fi
}

detect_threads() {
    if [[ "$THREADS" == "auto" ]]; then
        THREADS=$(nproc 2>/dev/null || echo 2)
        # Reserve 1 core for the management thread, minimum 1 worker
        local workers=$((THREADS - 1))
        [[ $workers -lt 1 ]] && workers=1
        THREADS=$workers
        log_info "Auto-detected CPU threads: $THREADS workers"
    fi
}

# ============================================================================
# INSTALLATION
# ============================================================================

install_suricata_debian() {
    log_step "Adding Suricata PPA / repository..."

    if [[ "$OS_ID" == "ubuntu" ]]; then
        apt-get update -qq
        apt-get install -y -qq software-properties-common
        add-apt-repository -y ppa:oisf/suricata-stable
        apt-get update -qq
    else
        # Debian
        apt-get update -qq
        apt-get install -y -qq gnupg2 apt-transport-https
    fi

    log_step "Installing Suricata and dependencies..."
    apt-get install -y -qq suricata suricata-update jq

    log_info "Suricata installed: $(suricata --build-info | head -1)"
}

install_suricata_rhel() {
    log_step "Installing EPEL and Suricata repository..."

    if [[ "$OS_ID" != "fedora" ]]; then
        dnf install -y -q epel-release
    fi

    dnf install -y -q dnf-plugins-core
    dnf copr enable -y @oisf/suricata-7.0 2>/dev/null || true

    log_step "Installing Suricata and dependencies..."
    dnf install -y -q suricata jq

    # Install suricata-update if not bundled
    if ! command -v suricata-update &>/dev/null; then
        dnf install -y -q python3-pip
        pip3 install suricata-update
    fi

    log_info "Suricata installed: $(suricata --build-info | head -1)"
}

install_suricata() {
    if command -v suricata &>/dev/null; then
        log_warn "Suricata is already installed: $(suricata -V 2>&1 | head -1)"
        log_info "Proceeding with configuration..."
        return 0
    fi

    case "$OS_FAMILY" in
        debian) install_suricata_debian ;;
        rhel)   install_suricata_rhel ;;
    esac
}

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

backup_config() {
    if [[ -f "$SURICATA_CONF" ]]; then
        local backup="${SURICATA_CONF}.bak.$(date +%Y%m%d%H%M%S)"
        cp "$SURICATA_CONF" "$backup"
        log_info "Backed up existing config to $backup"
    fi
}

configure_suricata() {
    log_step "Configuring Suricata..."

    backup_config

    # Set HOME_NET
    if grep -q 'HOME_NET:' "$SURICATA_CONF"; then
        sed -i "s|HOME_NET:.*|HOME_NET: \"$HOME_NET\"|" "$SURICATA_CONF"
        log_info "Set HOME_NET: $HOME_NET"
    fi

    # Set EXTERNAL_NET
    if grep -q 'EXTERNAL_NET:' "$SURICATA_CONF"; then
        sed -i 's|EXTERNAL_NET:.*|EXTERNAL_NET: "!\$HOME_NET"|' "$SURICATA_CONF"
    fi

    # Set interface
    sed -i "s|interface: eth0|interface: $IFACE|g" "$SURICATA_CONF" 2>/dev/null || true

    # Configure af-packet interface
    if grep -q 'af-packet:' "$SURICATA_CONF"; then
        # Update the first af-packet interface entry
        sed -i "/af-packet:/,/- interface:/{s|- interface:.*|- interface: $IFACE|}" "$SURICATA_CONF" 2>/dev/null || true
    fi

    # Set threading
    if [[ "$THREADS" -gt 1 ]]; then
        sed -i "s|#\?  threads:.*|  threads: $THREADS|" "$SURICATA_CONF" 2>/dev/null || true
        log_info "Set worker threads: $THREADS"
    fi

    # Enable community-id for EVE (useful for correlation)
    sed -i 's|# community-id:.*|community-id: true|' "$SURICATA_CONF" 2>/dev/null || true
    sed -i 's|community-id: false|community-id: true|' "$SURICATA_CONF" 2>/dev/null || true

    # Configure IPS mode if requested
    if [[ "$MODE" == "ips" ]]; then
        log_step "Configuring IPS (inline) mode..."

        # Set default action to drop for IPS
        if grep -q '# default-rule-path' "$SURICATA_CONF" || grep -q 'default-rule-path' "$SURICATA_CONF"; then
            log_info "IPS mode: rules with 'alert' will log, rules with 'drop' will block"
        fi

        # Configure NFQ for inline mode
        cat >> "$SURICATA_CONF" <<'NFQEOF'

# NFQ inline mode (IPS)
nfq:
  mode: accept
  fail-open: yes
NFQEOF

        log_info "IPS mode configured (NFQ with fail-open)"
        log_warn "You must configure iptables/nftables to send traffic to NFQUEUE"
        log_warn "Example: iptables -I FORWARD -j NFQUEUE --queue-num 0"
    fi

    # Create log directory
    mkdir -p "$(dirname "$EVE_LOG")"
    chown suricata:suricata "$(dirname "$EVE_LOG")" 2>/dev/null || true

    log_info "Configuration written to $SURICATA_CONF"
}

# ============================================================================
# RULE MANAGEMENT
# ============================================================================

configure_rules() {
    if [[ "$SKIP_RULES" == true ]]; then
        log_info "Skipping rule download (--skip-rules)"
        return 0
    fi

    log_step "Configuring rule sources..."

    # Parse comma-separated rule sources
    IFS=',' read -ra SOURCES <<< "$RULE_SOURCES"

    for source in "${SOURCES[@]}"; do
        source=$(echo "$source" | xargs)  # trim whitespace
        log_info "Enabling rule source: $source"

        case "$source" in
            et/open)
                suricata-update enable-source et/open 2>/dev/null || true
                ;;
            et/pro)
                if [[ -z "${OINKCODE:-}" ]]; then
                    log_warn "ET Pro requires OINKCODE environment variable — skipping"
                    continue
                fi
                suricata-update enable-source "et/pro secret-code=$OINKCODE" 2>/dev/null || true
                ;;
            oisf/trafficid)
                suricata-update enable-source oisf/trafficid 2>/dev/null || true
                ;;
            ptresearch/attackdetection)
                suricata-update enable-source ptresearch/attackdetection 2>/dev/null || true
                ;;
            sslbl/ssl-fp-blacklist)
                suricata-update enable-source sslbl/ssl-fp-blacklist 2>/dev/null || true
                ;;
            sslbl/ja3-fingerprints)
                suricata-update enable-source sslbl/ja3-fingerprints 2>/dev/null || true
                ;;
            etnetera/aggressive)
                suricata-update enable-source etnetera/aggressive 2>/dev/null || true
                ;;
            tgreen/hunting)
                suricata-update enable-source tgreen/hunting 2>/dev/null || true
                ;;
            malsilo/win-malware)
                suricata-update enable-source malsilo/win-malware 2>/dev/null || true
                ;;
            stamus/lateral)
                suricata-update enable-source stamus/lateral 2>/dev/null || true
                ;;
            *)
                log_warn "Unknown rule source: $source — trying as custom source"
                suricata-update enable-source "$source" 2>/dev/null || \
                    log_warn "Failed to enable source: $source"
                ;;
        esac
    done

    log_step "Downloading and installing rules..."
    suricata-update update

    # Show enabled sources
    log_info "Enabled rule sources:"
    suricata-update list-sources --enabled 2>/dev/null || true

    # Count installed rules
    local rule_count
    rule_count=$(suricata-update list-enabled-sources 2>/dev/null | wc -l || echo "unknown")
    local total_rules
    total_rules=$(grep -c '^\(alert\|drop\|reject\|pass\)' /var/lib/suricata/rules/suricata.rules 2>/dev/null || true)
    log_info "Total rules installed: $total_rules"
}

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

    cat > /etc/cron.d/suricata-update <<'CRONEOF'
# Update Suricata rules daily at 03:00 and reload
0 3 * * * root suricata-update update && suricatasc -c reload-rules 2>/dev/null
CRONEOF

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

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

configure_service() {
    log_step "Configuring Suricata systemd service..."

    # Ensure the service file uses the correct interface
    local service_file="/etc/systemd/system/suricata.service"
    local default_service="/lib/systemd/system/suricata.service"

    if [[ -f "$default_service" ]] && ! [[ -f "$service_file" ]]; then
        # Check if default service needs interface override
        if grep -q '%i' "$default_service" || grep -q 'af-packet' "$default_service"; then
            log_info "Service file uses af-packet config from suricata.yaml"
        fi
    fi

    # Enable and start
    systemctl daemon-reload
    systemctl enable suricata
    systemctl restart suricata

    # Wait for startup
    sleep 3

    if systemctl is-active --quiet suricata; then
        log_info "Suricata is running"
    else
        log_error "Suricata failed to start — check: journalctl -u suricata"
        systemctl status suricata --no-pager || true
    fi
}

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

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

    # Test config
    log_info "Testing configuration..."
    if suricata -T -c "$SURICATA_CONF" 2>&1 | tail -1 | grep -q 'Configuration provided was successfully loaded'; then
        log_info "Configuration is valid"
    else
        log_warn "Configuration test produced warnings (may still be functional)"
        suricata -T -c "$SURICATA_CONF" 2>&1 | tail -5
    fi

    # Check EVE log
    if [[ -f "$EVE_LOG" ]]; then
        log_info "EVE log exists: $EVE_LOG"
        local line_count
        line_count=$(wc -l < "$EVE_LOG" 2>/dev/null || echo 0)
        log_info "EVE log entries: $line_count"
    else
        log_info "EVE log will be created once traffic is processed"
    fi

    # Show status
    echo ""
    log_info "===== Installation Summary ====="
    log_info "Mode:        $MODE"
    log_info "Interface:   $IFACE"
    log_info "Config:      $SURICATA_CONF"
    log_info "EVE log:     $EVE_LOG"
    log_info "Rules:       $RULE_SOURCES"
    log_info "Threads:     $THREADS"
    log_info "HOME_NET:    $HOME_NET"
    echo ""
    suricata -V 2>&1 || true
}

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

uninstall_suricata() {
    log_step "Uninstalling Suricata..."

    systemctl stop suricata 2>/dev/null || true
    systemctl disable suricata 2>/dev/null || true

    case "$OS_FAMILY" in
        debian)
            apt-get remove --purge -y suricata suricata-update 2>/dev/null || true
            apt-get autoremove -y 2>/dev/null || true
            ;;
        rhel)
            dnf remove -y suricata 2>/dev/null || true
            ;;
    esac

    rm -f /etc/cron.d/suricata-update

    log_info "Suricata removed"
    log_info "Configuration files in /etc/suricata/ were left intact"
    log_info "Log files in /var/log/suricata/ were left intact"
    log_info "Rule files in /var/lib/suricata/ were left intact"
    log_info "Remove these directories 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 "Interface:   $IFACE"
    log_info "Mode:        $MODE"
    log_info "Rules:       $RULE_SOURCES"
    log_info "HOME_NET:    $HOME_NET"
    log_info "Threads:     $THREADS"
    log_info "EVE log:     $EVE_LOG"
    log_info "Config:      $SURICATA_CONF"
    echo ""
    log_info "Actions that would be performed:"
    echo "  1. Install Suricata via $( [[ $OS_FAMILY == debian ]] && echo 'apt' || echo 'dnf')"
    echo "  2. Configure $SURICATA_CONF"
    echo "  3. Set HOME_NET: $HOME_NET"
    echo "  4. Set interface: $IFACE"
    echo "  5. Set worker threads: $THREADS"
    echo "  6. Enable rule sources: $RULE_SOURCES"
    echo "  7. Download and install rules via suricata-update"
    echo "  8. Create daily rule update cron job"
    echo "  9. Enable and start suricata systemd service"
    if [[ "$MODE" == "ips" ]]; then
        echo "  10. Configure NFQ inline mode (IPS)"
    fi
    echo ""
}

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

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

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

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

    echo ""
    log_info "===== Suricata IDS/IPS Installer ====="
    echo ""

    install_suricata
    configure_suricata
    configure_rules
    setup_rule_update_cron
    configure_service
    validate_installation

    echo ""
    log_info "===== Installation Complete ====="
    log_info "Monitor alerts: tail -f $EVE_LOG | jq 'select(.event_type==\"alert\")'"
    log_info "View stats:     suricatasc -c dump-counters"
    log_info "Reload rules:   suricatasc -c reload-rules"
    echo ""
}

main "$@"