#!/usr/bin/env bash

######################################################################################
#### firewall-rule-diff.sh — Detect firewall rule drift against a saved baseline  ####
#### Supports UFW, iptables, and nftables. Saves snapshots, diffs against         ####
#### baseline, exports Prometheus metrics via textfile collector.                 ####
#### Requires: bash 4+, diff, coreutils                                           ####
####                                                                              ####
#### Author: Phil Connor                                                          ####
#### Contact: contact@mylinux.work                                                ####
#### License: MIT                                                                 ####
#### Version 1.00                                                                 ####
####                                                                              ####
#### Usage:                                                                       ####
####   sudo ./firewall-rule-diff.sh --save                                        ####
####   sudo ./firewall-rule-diff.sh --check                                       ####
####                                                                              ####
#### See --help for all options.                                                  ####
######################################################################################

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────
MODE=""  # save or check
BACKEND=""  # auto-detect: ufw, iptables, nftables
BASELINE_DIR="/etc/firewall-baseline"
MAX_AGE_DAYS=30
TEXTFILE_MODE=false
PROM_FILE="/var/lib/node_exporter/firewall_drift.prom"
OUTPUT_FORMAT="${OUTPUT_FORMAT:-text}"
JUNIT_FILE="${JUNIT_FILE:-firewall-drift-results.xml}"
VERBOSE="${VERBOSE:-false}"
COLOR="${COLOR:-auto}"

# ── State ─────────────────────────────────────────────────────────────
PASS=0
FAIL=0
WARN=0
TOTAL=0
RESULTS=()
START_TIME=""
RULES_ADDED=0
RULES_REMOVED=0
RULES_TOTAL=0
DRIFT_DETECTED=0
BASELINE_AGE=0
DETECTED_BACKEND=""

# ── Colors ────────────────────────────────────────────────────────────
setup_colors() {
    if [[ "$COLOR" == "never" ]]; then
        RED="" GREEN="" YELLOW="" BLUE="" BOLD="" RESET=""
        return
    fi
    if [[ "$COLOR" == "always" ]] || [[ -t 1 ]]; then
        RED='\033[0;31m'
        GREEN='\033[0;32m'
        YELLOW='\033[0;33m'
        BLUE='\033[0;34m'
        BOLD='\033[1m'
        RESET='\033[0m'
    else
        RED="" GREEN="" YELLOW="" BLUE="" BOLD="" RESET=""
    fi
}

# ── Logging ───────────────────────────────────────────────────────────
log()     { echo -e "${BLUE}[INFO]${RESET} $*"; }
warn()    { echo -e "${YELLOW}[WARN]${RESET} $*" >&2; }
err()     { echo -e "${RED}[ERROR]${RESET} $*" >&2; }
verbose() { if [[ "$VERBOSE" == "true" ]]; then echo -e "${BLUE}[DEBUG]${RESET} $*"; fi; }

# ── Test Result Recording ─────────────────────────────────────────────
record_pass() {
    local name="$1"
    local detail="${2:-}"
    ((PASS++)) || true
    ((TOTAL++)) || true
    RESULTS+=("PASS|${name}|${detail}")
    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then
        echo "ok ${TOTAL} - ${name}"
    elif [[ "$OUTPUT_FORMAT" == "text" ]]; then
        echo -e "  ${GREEN}✓${RESET} ${name}${detail:+ — ${detail}}"
    fi
}

record_fail() {
    local name="$1"
    local detail="${2:-}"
    ((FAIL++)) || true
    ((TOTAL++)) || true
    RESULTS+=("FAIL|${name}|${detail}")
    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then
        echo "not ok ${TOTAL} - ${name}"
        [[ -n "$detail" ]] && echo "  # ${detail}"
    elif [[ "$OUTPUT_FORMAT" == "text" ]]; then
        echo -e "  ${RED}✗${RESET} ${name}${detail:+ — ${detail}}"
    fi
}

record_warn() {
    local name="$1"
    local detail="${2:-}"
    ((WARN++)) || true
    ((TOTAL++)) || true
    RESULTS+=("WARN|${name}|${detail}")
    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then
        echo "ok ${TOTAL} - ${name} # SKIP ${detail}"
    elif [[ "$OUTPUT_FORMAT" == "text" ]]; then
        echo -e "  ${YELLOW}⊘${RESET} ${name}${detail:+ — ${detail}}"
    fi
}

# ── Help ──────────────────────────────────────────────────────────────
show_help() {
    cat <<'EOF'
Usage: firewall-rule-diff.sh [OPTIONS]

Detect firewall rule drift by comparing current state against a saved baseline.
Supports UFW, iptables, and nftables with auto-detection.

Modes:
  --save                Save current firewall rules as new baseline
  --check               Compare current rules against baseline (default)

Options:
  --backend BACKEND     Force backend: ufw, iptables, nftables (default: auto-detect)
  --baseline-dir PATH   Baseline storage directory (default: /etc/firewall-baseline/)
  --max-age DAYS        Warn if baseline older than N days (default: 30)
  --textfile            Write Prometheus metrics to textfile collector
  --prom-file PATH      Textfile path (default: /var/lib/node_exporter/firewall_drift.prom)
  --format FORMAT       Output: text (default), tap, junit
  --junit-file FILE     JUnit output path (default: firewall-drift-results.xml)
  --verbose             Show debug output
  --no-color            Disable colored output
  -h, --help            Show this help

Examples:
  sudo ./firewall-rule-diff.sh --save
  sudo ./firewall-rule-diff.sh --check
  sudo ./firewall-rule-diff.sh --check --textfile
  sudo ./firewall-rule-diff.sh --backend iptables --check
  sudo ./firewall-rule-diff.sh --check --max-age 7
EOF
    exit 0
}

# ── Parse Arguments ───────────────────────────────────────────────────
parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --save)           MODE="save"; shift ;;
            --check)          MODE="check"; shift ;;
            --backend)        BACKEND="$2"; shift 2 ;;
            --baseline-dir)   BASELINE_DIR="$2"; shift 2 ;;
            --max-age)        MAX_AGE_DAYS="$2"; shift 2 ;;
            --textfile)       TEXTFILE_MODE=true; shift ;;
            --prom-file)      PROM_FILE="$2"; shift 2 ;;
            --format)         OUTPUT_FORMAT="$2"; shift 2 ;;
            --junit-file)     JUNIT_FILE="$2"; shift 2 ;;
            --verbose)        VERBOSE=true; shift ;;
            --no-color)       COLOR="never"; shift ;;
            -h|--help)        show_help ;;
            *)                err "Unknown option: $1"; echo "Run with --help for usage."; exit 1 ;;
        esac
    done

    if [[ -z "$MODE" ]]; then
        MODE="check"
    fi
}

# ── Detect Backend ────────────────────────────────────────────────────
detect_backend() {
    if [[ -n "$BACKEND" ]]; then
        DETECTED_BACKEND="$BACKEND"
        verbose "Backend forced: ${DETECTED_BACKEND}"
        return
    fi

    if command -v ufw &>/dev/null && ufw status &>/dev/null; then
        local ufw_status
        ufw_status=$(ufw status 2>/dev/null | head -1)
        if [[ "$ufw_status" == *"active"* ]]; then
            DETECTED_BACKEND="ufw"
            verbose "Detected active UFW"
            return
        fi
    fi

    if command -v nft &>/dev/null; then
        local nft_rules
        nft_rules=$(nft list ruleset 2>/dev/null | wc -l)
        if [[ "$nft_rules" -gt 0 ]]; then
            DETECTED_BACKEND="nftables"
            verbose "Detected nftables with ${nft_rules} lines"
            return
        fi
    fi

    if command -v iptables-save &>/dev/null; then
        DETECTED_BACKEND="iptables"
        verbose "Detected iptables"
        return
    fi

    err "No supported firewall backend found (ufw, nftables, iptables)"
    exit 1
}

# ── Snapshot Functions ────────────────────────────────────────────────
snapshot_ufw() {
    local dir="$1"
    ufw status numbered > "${dir}/ufw-status.txt" 2>/dev/null || true
    ufw status verbose > "${dir}/ufw-verbose.txt" 2>/dev/null || true
    if [[ -f /etc/ufw/user.rules ]]; then
        cp /etc/ufw/user.rules "${dir}/user.rules"
    fi
    if [[ -f /etc/ufw/user6.rules ]]; then
        cp /etc/ufw/user6.rules "${dir}/user6.rules"
    fi
    # count rules from numbered output (skip header lines)
    RULES_TOTAL=$(grep -cE '^\[' "${dir}/ufw-status.txt" 2>/dev/null) || RULES_TOTAL=0
    verbose "UFW snapshot: ${RULES_TOTAL} rules"
}

snapshot_iptables() {
    local dir="$1"
    iptables-save > "${dir}/iptables-v4.rules" 2>/dev/null || true
    if command -v ip6tables-save &>/dev/null; then
        ip6tables-save > "${dir}/iptables-v6.rules" 2>/dev/null || true
    fi
    # count non-comment, non-empty lines
    RULES_TOTAL=$(grep -cvE '^(#|$|\*|COMMIT|:)' "${dir}/iptables-v4.rules" 2>/dev/null) || RULES_TOTAL=0
    if [[ -f "${dir}/iptables-v6.rules" ]]; then
        local v6_count
        v6_count=$(grep -cvE '^(#|$|\*|COMMIT|:)' "${dir}/iptables-v6.rules" 2>/dev/null) || v6_count=0
        RULES_TOTAL=$((RULES_TOTAL + v6_count))
    fi
    verbose "iptables snapshot: ${RULES_TOTAL} rules"
}

snapshot_nftables() {
    local dir="$1"
    nft list ruleset > "${dir}/nftables.rules" 2>/dev/null || true
    RULES_TOTAL=$(grep -cE '^\s+(rule|chain|table)' "${dir}/nftables.rules" 2>/dev/null) || RULES_TOTAL=0
    verbose "nftables snapshot: ${RULES_TOTAL} lines"
}

take_snapshot() {
    local dir="$1"
    case "$DETECTED_BACKEND" in
        ufw)      snapshot_ufw "$dir" ;;
        iptables) snapshot_iptables "$dir" ;;
        nftables) snapshot_nftables "$dir" ;;
    esac
    echo "$DETECTED_BACKEND" > "${dir}/backend.txt"
    date +%s > "${dir}/timestamp.txt"
    date -Is > "${dir}/timestamp-human.txt"
}

# ── Save Mode ─────────────────────────────────────────────────────────
do_save() {
    mkdir -p "${BASELINE_DIR}"
    local snapshot_dir="${BASELINE_DIR}/baseline"

    # clean up any previous baseline
    if [[ -d "$snapshot_dir" ]]; then
        local prev_ts
        prev_ts=$(cat "${snapshot_dir}/timestamp.txt" 2>/dev/null || echo "unknown")
        local archive_dir="${BASELINE_DIR}/archive-${prev_ts}"
        mv "$snapshot_dir" "$archive_dir" 2>/dev/null || rm -rf "$snapshot_dir"
        verbose "Archived previous baseline"
    fi

    mkdir -p "$snapshot_dir"

    verbose "Taking ${DETECTED_BACKEND} snapshot to ${snapshot_dir}"
    take_snapshot "$snapshot_dir"

    if [[ ! -f "${snapshot_dir}/backend.txt" ]]; then
        err "Snapshot failed — ${snapshot_dir}/backend.txt not created"
        exit 1
    fi

    if [[ "$OUTPUT_FORMAT" == "text" ]]; then
        echo -e "${BOLD}Firewall Rule Diff${RESET}"
        echo "Backend:   ${DETECTED_BACKEND}"
        echo "Baseline:  ${snapshot_dir}"
        echo "Time:      $(cat "${snapshot_dir}/timestamp-human.txt")"
        echo "Rules:     ${RULES_TOTAL}"
        echo ""
        echo -e "  ${GREEN}✓${RESET} Baseline saved — ${RULES_TOTAL} rules (${DETECTED_BACKEND})"
    elif [[ "$OUTPUT_FORMAT" == "tap" ]]; then
        echo "1..1"
        echo "ok 1 - Baseline saved (${RULES_TOTAL} rules, ${DETECTED_BACKEND})"
    fi
}

# ── Check Mode ────────────────────────────────────────────────────────
do_check() {
    START_TIME=$(date +%s)
    local baseline_dir="${BASELINE_DIR}/baseline"

    if [[ "$OUTPUT_FORMAT" == "text" ]]; then
        echo -e "${BOLD}Firewall Rule Diff${RESET}"
        echo "Backend:      ${DETECTED_BACKEND}"
        echo "Baseline:     ${baseline_dir}"
        echo "Time:         $(date -Is)"
        echo ""
    elif [[ "$OUTPUT_FORMAT" == "tap" ]]; then
        echo "TAP version 13"
    fi

    # check baseline exists
    if [[ ! -d "$baseline_dir" ]]; then
        record_fail "Baseline exists" "no baseline found — run with --save first"
        DRIFT_DETECTED=1
        print_summary
        return
    fi

    # check backend matches
    local baseline_backend
    baseline_backend=$(cat "${baseline_dir}/backend.txt" 2>/dev/null || echo "unknown")
    if [[ "$baseline_backend" != "$DETECTED_BACKEND" ]]; then
        record_fail "Backend match" "baseline uses ${baseline_backend}, current is ${DETECTED_BACKEND}"
        DRIFT_DETECTED=1
    else
        record_pass "Backend match" "${DETECTED_BACKEND}"
    fi

    # check baseline age
    local baseline_ts
    baseline_ts=$(cat "${baseline_dir}/timestamp.txt" 2>/dev/null || echo "0")
    local now
    now=$(date +%s)
    BASELINE_AGE=$((now - baseline_ts))
    local age_days=$((BASELINE_AGE / 86400))

    if [[ $age_days -gt $MAX_AGE_DAYS ]]; then
        record_warn "Baseline age" "${age_days} days old (threshold: ${MAX_AGE_DAYS})"
    else
        record_pass "Baseline age" "${age_days} days old"
    fi

    # take current snapshot to temp dir
    local tmp_dir
    tmp_dir=$(mktemp -d)
    trap 'rm -rf "'"$tmp_dir"'"' EXIT
    take_snapshot "$tmp_dir"

    if [[ "$OUTPUT_FORMAT" == "text" ]]; then
        echo ""
        echo -e "${BOLD}Rule Comparison${RESET}"
    fi

    # diff based on backend
    case "$DETECTED_BACKEND" in
        ufw)      diff_ufw "$baseline_dir" "$tmp_dir" ;;
        iptables) diff_iptables "$baseline_dir" "$tmp_dir" ;;
        nftables) diff_nftables "$baseline_dir" "$tmp_dir" ;;
    esac

    print_summary
}

# ── Diff Functions ────────────────────────────────────────────────────
diff_rules_file() {
    local label="$1"
    local baseline_file="$2"
    local current_file="$3"

    if [[ ! -f "$baseline_file" ]] && [[ ! -f "$current_file" ]]; then
        verbose "Both files missing for ${label} — skipping"
        return
    fi

    if [[ ! -f "$baseline_file" ]]; then
        record_fail "${label}" "file missing from baseline but present now"
        DRIFT_DETECTED=1
        return
    fi

    if [[ ! -f "$current_file" ]]; then
        record_fail "${label}" "file present in baseline but missing now"
        DRIFT_DETECTED=1
        return
    fi

    local diff_output
    diff_output=$(diff --unified=0 "$baseline_file" "$current_file" 2>/dev/null) || true

    if [[ -z "$diff_output" ]]; then
        record_pass "${label}" "no changes"
        return
    fi

    DRIFT_DETECTED=1

    local added removed
    added=$(echo "$diff_output" | grep -c '^+[^+]' 2>/dev/null) || added=0
    removed=$(echo "$diff_output" | grep -c '^-[^-]' 2>/dev/null) || removed=0

    RULES_ADDED=$((RULES_ADDED + added))
    RULES_REMOVED=$((RULES_REMOVED + removed))

    record_fail "${label}" "${added} added, ${removed} removed"

    if [[ "$VERBOSE" == "true" || "$OUTPUT_FORMAT" == "text" ]]; then
        # show the actual diff lines (limit to 20 lines)
        local count=0
        while IFS= read -r line; do
            if [[ "$line" == +* && "$line" != +++* ]]; then
                echo -e "    ${GREEN}${line}${RESET}"
                ((count++)) || true
            elif [[ "$line" == -* && "$line" != ---* ]]; then
                echo -e "    ${RED}${line}${RESET}"
                ((count++)) || true
            fi
            [[ $count -ge 20 ]] && { echo "    ... (truncated)"; break; }
        done <<< "$diff_output"
    fi
}

diff_ufw() {
    local baseline="$1"
    local current="$2"

    diff_rules_file "UFW status" "${baseline}/ufw-status.txt" "${current}/ufw-status.txt"
    diff_rules_file "UFW IPv4 rules" "${baseline}/user.rules" "${current}/user.rules"
    diff_rules_file "UFW IPv6 rules" "${baseline}/user6.rules" "${current}/user6.rules"

    # rule count comparison
    local baseline_count current_count
    baseline_count=$(grep -cE '^\[' "${baseline}/ufw-status.txt" 2>/dev/null) || baseline_count=0
    current_count=$(grep -cE '^\[' "${current}/ufw-status.txt" 2>/dev/null) || current_count=0

    if [[ $baseline_count -ne $current_count ]]; then
        record_fail "Rule count" "baseline: ${baseline_count}, current: ${current_count}"
        DRIFT_DETECTED=1
    else
        record_pass "Rule count" "${current_count} rules"
    fi
}

diff_iptables() {
    local baseline="$1"
    local current="$2"

    diff_rules_file "iptables IPv4 rules" "${baseline}/iptables-v4.rules" "${current}/iptables-v4.rules"
    diff_rules_file "iptables IPv6 rules" "${baseline}/iptables-v6.rules" "${current}/iptables-v6.rules"

    # chain count comparison
    local baseline_chains current_chains
    baseline_chains=$(grep -cE '^:' "${baseline}/iptables-v4.rules" 2>/dev/null) || baseline_chains=0
    current_chains=$(grep -cE '^:' "${current}/iptables-v4.rules" 2>/dev/null) || current_chains=0

    if [[ $baseline_chains -ne $current_chains ]]; then
        record_fail "Chain count (IPv4)" "baseline: ${baseline_chains}, current: ${current_chains}"
        DRIFT_DETECTED=1
    else
        record_pass "Chain count (IPv4)" "${current_chains} chains"
    fi
}

diff_nftables() {
    local baseline="$1"
    local current="$2"

    diff_rules_file "nftables ruleset" "${baseline}/nftables.rules" "${current}/nftables.rules"

    # table count comparison
    local baseline_tables current_tables
    baseline_tables=$(grep -c '^table' "${baseline}/nftables.rules" 2>/dev/null) || baseline_tables=0
    current_tables=$(grep -c '^table' "${current}/nftables.rules" 2>/dev/null) || current_tables=0

    if [[ $baseline_tables -ne $current_tables ]]; then
        record_fail "Table count" "baseline: ${baseline_tables}, current: ${current_tables}"
        DRIFT_DETECTED=1
    else
        record_pass "Table count" "${current_tables} tables"
    fi
}

# ── Summary ───────────────────────────────────────────────────────────
print_summary() {
    local end_time
    end_time=$(date +%s)
    local elapsed=$((end_time - START_TIME))

    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then
        echo "1..${TOTAL}"
    elif [[ "$OUTPUT_FORMAT" == "text" ]]; then
        echo ""
        echo "────────────────────────────────────────"
        echo -e "${BOLD}Summary${RESET}  ${DETECTED_BACKEND}"
        echo -e "  ${PASS} passed  ${FAIL} failed  ${WARN} skipped  (${elapsed}s)"
        if [[ $DRIFT_DETECTED -eq 1 ]]; then
            echo -e "  Rules added: ${RULES_ADDED}  removed: ${RULES_REMOVED}"
            echo -e "  ${RED}Drift detected.${RESET}"
        else
            echo -e "  ${GREEN}No drift detected.${RESET}"
        fi
        echo "────────────────────────────────────────"
    fi

    if [[ "$OUTPUT_FORMAT" == "junit" ]]; then
        write_junit
    fi

    if [[ "$TEXTFILE_MODE" == "true" ]]; then
        write_prometheus
    fi
}

# ── JUnit Output ──────────────────────────────────────────────────────
write_junit() {
    local end_time
    end_time=$(date +%s)
    local elapsed=$((end_time - START_TIME))

    {
        echo '<?xml version="1.0" encoding="UTF-8"?>'
        echo "<testsuites tests=\"${TOTAL}\" failures=\"${FAIL}\" time=\"${elapsed}\">"
        echo "  <testsuite name=\"firewall-rule-diff\" tests=\"${TOTAL}\" failures=\"${FAIL}\" skipped=\"${WARN}\" time=\"${elapsed}\">"

        for result in "${RESULTS[@]}"; do
            local status name detail
            status=$(echo "$result" | cut -d'|' -f1)
            name=$(echo "$result" | cut -d'|' -f2)
            detail=$(echo "$result" | cut -d'|' -f3)

            # escape XML
            name="${name//&/&amp;}"
            name="${name//</&lt;}"
            name="${name//>/&gt;}"
            detail="${detail//&/&amp;}"
            detail="${detail//</&lt;}"
            detail="${detail//>/&gt;}"

            echo "    <testcase name=\"${name}\">"
            if [[ "$status" == "FAIL" ]]; then
                echo "      <failure message=\"${detail}\"></failure>"
            elif [[ "$status" == "WARN" ]]; then
                echo "      <skipped message=\"${detail}\"></skipped>"
            fi
            echo "    </testcase>"
        done

        echo "  </testsuite>"
        echo "</testsuites>"
    } > "$JUNIT_FILE"

    verbose "JUnit report written to ${JUNIT_FILE}"
}

# ── Prometheus Output ─────────────────────────────────────────────────
write_prometheus() {
    local prom_dir
    prom_dir=$(dirname "$PROM_FILE")
    if [[ ! -d "$prom_dir" ]]; then
        warn "Prometheus textfile directory does not exist: ${prom_dir}"
        return
    fi

    local tmp_file="${PROM_FILE}.$$"
    {
        echo "# HELP firewall_drift_detected Whether firewall rules differ from baseline"
        echo "# TYPE firewall_drift_detected gauge"
        echo "firewall_drift_detected ${DRIFT_DETECTED}"
        echo "# HELP firewall_rules_added Rules added since baseline"
        echo "# TYPE firewall_rules_added gauge"
        echo "firewall_rules_added ${RULES_ADDED}"
        echo "# HELP firewall_rules_removed Rules removed since baseline"
        echo "# TYPE firewall_rules_removed gauge"
        echo "firewall_rules_removed ${RULES_REMOVED}"
        echo "# HELP firewall_rules_total Current total firewall rules"
        echo "# TYPE firewall_rules_total gauge"
        echo "firewall_rules_total ${RULES_TOTAL}"
        echo "# HELP firewall_baseline_age_seconds Seconds since baseline was saved"
        echo "# TYPE firewall_baseline_age_seconds gauge"
        echo "firewall_baseline_age_seconds ${BASELINE_AGE}"
        echo "# HELP firewall_scan_timestamp Unix timestamp of last scan"
        echo "# TYPE firewall_scan_timestamp gauge"
        echo "firewall_scan_timestamp $(date +%s)"
        echo "# HELP firewall_backend Active firewall backend"
        echo "# TYPE firewall_backend gauge"
        echo "firewall_backend{backend=\"${DETECTED_BACKEND}\"} 1"
    } > "$tmp_file"

    mv "$tmp_file" "$PROM_FILE"
    verbose "Prometheus metrics written to ${PROM_FILE}"
}

# ── Main ──────────────────────────────────────────────────────────────
main() {
    setup_colors
    parse_args "$@"
    setup_colors  # re-apply after --no-color

    if [[ $EUID -ne 0 ]]; then
        err "This script must be run as root."
        exit 1
    fi

    detect_backend

    case "$MODE" in
        save)  do_save ;;
        check) do_check ;;
    esac

    if [[ $FAIL -gt 0 ]]; then
        exit 1
    fi
    exit 0
}

main "$@"