#!/usr/bin/env bash

set -u

SNORT_BIN="/usr/local/bin/snort"
SNORT_CONF="/usr/local/etc/snort/snort.lua"
INTERFACE="ens33"

LOG_DIR="/var/log/snort/alerts"
ALERT_FILE="$LOG_DIR/alert_fast.txt"
RUN_DIR="/var/run/snort"
PID_FILE="$RUN_DIR/snort_live.pid"
STDOUT_LOG="$LOG_DIR/snort_live_stdout.log"
STDERR_LOG="$LOG_DIR/snort_live_stderr.log"

LOG_PCAP_BASE="/var/log/snort/pcap_results"
PCAP_BASE="/mnt/hgfs/pcapdataset"

REPORT_SCRIPT="/usr/local/bin/nids-report.py"

mkdir -p "$LOG_DIR"
mkdir -p "$RUN_DIR"
mkdir -p "$LOG_PCAP_BASE"

start_live() {
    echo "[+] Starting Snort on interface $INTERFACE"

    if [ -f "$PID_FILE" ]; then
        OLD_PID=$(cat "$PID_FILE" 2>/dev/null || true)
        if [ -n "${OLD_PID:-}" ] && ps -p "$OLD_PID" > /dev/null 2>&1; then
            echo "[!] Snort is already running with PID $OLD_PID"
            exit 1
        else
            rm -f "$PID_FILE"
        fi
    fi

    if [ ! -x "$SNORT_BIN" ]; then
        echo "[!] Snort binary not found: $SNORT_BIN"
        exit 1
    fi

    if [ ! -f "$SNORT_CONF" ]; then
        echo "[!] Snort config not found: $SNORT_CONF"
        exit 1
    fi

    echo "[*] Testing config..."
    "$SNORT_BIN" -c "$SNORT_CONF" -T > /tmp/snort_test.log 2>&1
    if [ $? -ne 0 ]; then
        echo "[!] Config test failed"
        cat /tmp/snort_test.log
        exit 1
    fi

    nohup "$SNORT_BIN" \
        -c "$SNORT_CONF" \
        -i "$INTERFACE" \
        -A alert_fast \
        -l "$LOG_DIR" \
        > "$STDOUT_LOG" 2> "$STDERR_LOG" &

    SNORT_PID=$!
    sleep 2

    if ps -p "$SNORT_PID" > /dev/null 2>&1; then
        echo "$SNORT_PID" > "$PID_FILE"
        echo "[+] Snort started with PID $SNORT_PID"
        echo "[*] Alerts: $ALERT_FILE"
        echo "[*] Stdout: $STDOUT_LOG"
        echo "[*] Stderr: $STDERR_LOG"
    else
        echo "[!] Snort failed to start"
        echo "----- STDERR -----"
        cat "$STDERR_LOG" 2>/dev/null || true
        echo "----- STDOUT -----"
        cat "$STDOUT_LOG" 2>/dev/null || true
        exit 1
    fi
}

stop_live() {
    echo "[+] Stopping Snort"

    if [ ! -f "$PID_FILE" ]; then
        echo "[!] No PID file found. Snort may not be running."
        exit 1
    fi

    SNORT_PID=$(cat "$PID_FILE" 2>/dev/null || true)

    if [ -z "${SNORT_PID:-}" ]; then
        echo "[!] PID file is empty"
        rm -f "$PID_FILE"
        exit 1
    fi

    if ps -p "$SNORT_PID" > /dev/null 2>&1; then
        kill "$SNORT_PID"
        sleep 2

        if ps -p "$SNORT_PID" > /dev/null 2>&1; then
            echo "[!] Snort did not stop, killing forcefully..."
            kill -9 "$SNORT_PID"
        fi

        echo "[+] Snort stopped"
    else
        echo "[!] Process $SNORT_PID is not running"
    fi

    rm -f "$PID_FILE"
    run_report
}

run_pcap() {
    SESSION_NAME="${1:-}"
    PCAP_FILE="${2:-}"

    if [ -z "$SESSION_NAME" ] || [ -z "$PCAP_FILE" ]; then
        echo "Usage: sudo nids-snort.sh pcap <tuesday|wednesday|friday> <path/to/file.pcap>"
        exit 1
    fi

    if [ ! -f "$PCAP_FILE" ]; then
        echo "[!] PCAP file not found: $PCAP_FILE"
        exit 1
    fi

    SESSION_LOWER=$(echo "$SESSION_NAME" | tr '[:upper:]' '[:lower:]')
    PCAP_LOG_DIR="$LOG_PCAP_BASE/$SESSION_LOWER"
    PCAP_STDOUT="$PCAP_LOG_DIR/snort_stdout.log"
    PCAP_STDERR="$PCAP_LOG_DIR/snort_stderr.log"

    mkdir -p "$PCAP_LOG_DIR"

    PCAP_PID=""

    # Ctrl+C stops Snort and starts report; second Ctrl+C skips report and exits
    trap '
        echo ""
        echo "[!] Interrupted, stopping Snort..."
        if [ -n "$PCAP_PID" ]; then
            kill "$PCAP_PID" 2>/dev/null || true
            wait "$PCAP_PID" 2>/dev/null || true
        fi
        ALERT_COUNT=$(wc -l < "$PCAP_LOG_DIR/alert_fast.txt" 2>/dev/null || echo 0)
        echo "[*] Alerts captured: $ALERT_COUNT"
        echo "[*] Generating report now, press Ctrl+C again to skip"

        trap "echo \"\"; echo \"[!] Report skipped. Exiting.\"; exit 0" INT

        python3 "$REPORT_SCRIPT"
        echo "[+] Report ready: /home/osboxes/nids-reports/report_all.html"

        trap - INT
        exit 0
    ' INT

    echo "[+] Starting Snort in PCAP mode"
    echo "[*] Session  : $SESSION_LOWER"
    echo "[*] PCAP file: $PCAP_FILE"
    echo "[*] Alert log: $PCAP_LOG_DIR/alert_fast.txt"
    echo "[*] Ctrl+C to stop and generate report"

    nohup "$SNORT_BIN" \
        -c "$SNORT_CONF" \
        -r "$PCAP_FILE" \
        -A alert_fast \
        -l "$PCAP_LOG_DIR" \
        > "$PCAP_STDOUT" 2> "$PCAP_STDERR" &

    PCAP_PID=$!
    echo "[*] Snort running in background with PID $PCAP_PID"
    echo "[*] Waiting for it to finish..."
    wait $PCAP_PID
    EXIT_CODE=$?

    trap - INT

    if [ $EXIT_CODE -eq 0 ]; then
        echo "[+] Snort finished: $SESSION_LOWER"
    else
        echo "[!] Snort exited with code $EXIT_CODE"
        echo "----- STDERR -----"
        cat "$PCAP_STDERR" 2>/dev/null || true
    fi

    ALERT_COUNT=$(wc -l < "$PCAP_LOG_DIR/alert_fast.txt" 2>/dev/null || echo 0)
    echo "[*] Alerts generated: $ALERT_COUNT"
}

run_all() {
    echo ""
    echo "[*] Running all 3 PCAP sessions then generating report"
    echo "[*] Ctrl+C during PCAP: stop + generate report"
    echo "[*] Ctrl+C during report: skip report + exit"
    echo ""

    echo "[*] Step 1/4 - PCAP: Tuesday (SSH Brute Force)"
    run_pcap "tuesday" "$PCAP_BASE/Tuesday-WorkingHours.pcap"

    echo ""
    echo "[*] Step 2/4 - PCAP: Wednesday (DoS/DDoS)"
    run_pcap "wednesday" "$PCAP_BASE/Wednesday-workingHours.pcap"

    echo ""
    echo "[*] Step 3/4 - PCAP: Friday (Port Scan)"
    run_pcap "friday" "$PCAP_BASE/Friday-WorkingHours.pcap"

    echo ""
    echo "[*] Step 4/4 - Generating HTML report"
    run_report

    echo ""
    echo "[+] All sessions done."
    echo "[+] Report: /home/osboxes/nids-reports/report_all.html"
}

run_report() {
    if [ -f "$REPORT_SCRIPT" ]; then
        echo "[*] Generating HTML report..."
        python3 "$REPORT_SCRIPT"
        echo "[+] Report ready: /home/osboxes/nids-reports/report_all.html"
    else
        echo "[!] Report script not found: $REPORT_SCRIPT"
    fi
}

status_live() {
    if [ -f "$PID_FILE" ]; then
        SNORT_PID=$(cat "$PID_FILE" 2>/dev/null || true)
        if [ -n "${SNORT_PID:-}" ] && ps -p "$SNORT_PID" > /dev/null 2>&1; then
            echo "[+] Snort is running with PID $SNORT_PID"
            exit 0
        fi
    fi
    echo "[!] Snort is not running"
    exit 1
}

tail_alerts() {
    SESSION="${1:-live}"
    SESSION_LOWER=$(echo "$SESSION" | tr '[:upper:]' '[:lower:]')

    if [ "$SESSION_LOWER" = "live" ]; then
        TARGET="$ALERT_FILE"
    else
        TARGET="$LOG_PCAP_BASE/$SESSION_LOWER/alert_fast.txt"
    fi

    if [ -f "$TARGET" ]; then
        echo "[*] Watching: $TARGET"
        tail -f "$TARGET"
    else
        echo "[!] Alert file not found: $TARGET"
        exit 1
    fi
}

show_logs() {
    echo "===== LIVE STDERR ====="
    cat "$STDERR_LOG" 2>/dev/null || true
    echo
    echo "===== LIVE STDOUT ====="
    cat "$STDOUT_LOG" 2>/dev/null || true
}

case "${1:-}" in
    live)
        start_live
        ;;
    stop)
        stop_live
        ;;
    pcap)
        run_pcap "${2:-}" "${3:-}"
        ;;
    all)
        run_all
        ;;
    report)
        run_report
        ;;
    status)
        status_live
        ;;
    tail)
        tail_alerts "${2:-live}"
        ;;
    logs)
        show_logs
        ;;
    *)
        echo ""
        echo "Usage: sudo nids-snort.sh {live|stop|pcap|all|report|status|tail|logs}"
        echo ""
        echo "  live                          Start Snort live monitoring"
        echo "  stop                          Stop Snort and generate report"
        echo "  pcap <session> <file.pcap>    Replay a single PCAP file"
        echo "  all                           Run all 3 PCAP sessions and report"
        echo "  report                        Generate HTML report"
        echo "  status                        Check if Snort is running"
        echo "  tail [live|tuesday|...]       Watch alert log in real time"
        echo "  logs                          Show Snort stdout/stderr"
        echo ""
        exit 1
        ;;
esac