#!/usr/bin/env bash
# test-firewall-scenarios.sh — 3-VM QEMU topology for firewall validation.
#
# Topology:
#   ┌─────────────┐     ┌─────────────┐     ┌─────────────┐
#   │   VM-A      │     │  Router-VM  │     │   VM-B      │
#   │ (internal)  │◄───►│  2 NICs     │◄───►│ (external)  │
#   │ 10.0.0.2/24 │     │ eth0: 10.0.0.1  │     │ eth1: dyn    │
#   └─────────────┘     │ eth1: 10.0.2.100│     └─────────────┘
#                       └─────────────┘
#
# Router-VM runs redbear-firewall-check to exercise the netfilter control
# plane. pcaps are captured on each NIC via QEMU filter-dump.
#
# Usage:
#   ./test-firewall-scenarios.sh [--check] [REDBEAR_CONFIG]
#
#   REDBEAR_CONFIG defaults to redbear-mini (text-only, smaller ISO).
#   --check: run in automated check mode using expect.
#
# Artifacts are saved to artifacts/firewall-<date>/.

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
ROOT_DIR="$(cd "$SCRIPT_DIR/../.." && pwd)"
ISO_CONFIG="${REDBEAR_CONFIG:-redbear-mini}"
ARCH="${ARCH:-$(uname -m)}"
TIMESTAMP="$(date +%Y%m%d-%H%M%S)"
ARTIFACT_DIR="$ROOT_DIR/artifacts/firewall-$TIMESTAMP"

declare -a vm_pids=()
declare -a sock_paths=()
cleanup_done=0

# ──────────────────────────────────────────────────────────────────────────
# Helpers
# ──────────────────────────────────────────────────────────────────────────

find_uefi_firmware() {
    local candidates=(
        "/usr/share/ovmf/x64/OVMF.4m.fd"
        "/usr/share/OVMF/x64/OVMF.4m.fd"
        "/usr/share/ovmf/x64/OVMF_CODE.4m.fd"
        "/usr/share/OVMF/x64/OVMF_CODE.4m.fd"
        "/usr/share/ovmf/OVMF.fd"
        "/usr/share/OVMF/OVMF_CODE.fd"
        "/usr/share/qemu/edk2-x86_64-code.fd"
    )
    for path in "${candidates[@]}"; do
        if [[ -f "$path" ]]; then
            printf '%s\n' "$path"
            return 0
        fi
    done
    return 1
}

safe_kill_vm() {
    local pid="$1"
    if kill -0 "$pid" 2>/dev/null; then
        kill "$pid" 2>/dev/null || true
        # Wait up to 5 seconds for graceful shutdown
        local waited=0
        while kill -0 "$pid" 2>/dev/null && [[ $waited -lt 10 ]]; do
            sleep 0.5
            waited=$((waited + 1))
        done
        # Force kill if still alive
        kill -9 "$pid" 2>/dev/null || true
        wait "$pid" 2>/dev/null || true
    fi
}

cleanup() {
    if [[ $cleanup_done -eq 1 ]]; then
        return
    fi
    cleanup_done=1
    echo "=== Cleaning up VMs ==="
    for pid in "${vm_pids[@]}"; do
        safe_kill_vm "$pid"
    done
    for sock in "${sock_paths[@]}"; do
        rm -f "$sock"
    done
    echo "VMs shut down."
    echo "Artifacts: $ARTIFACT_DIR"
}
trap cleanup EXIT INT TERM

usage() {
    cat <<'USAGE'
Usage: test-firewall-scenarios.sh [--check] [REDBEAR_CONFIG]

3-VM QEMU topology for firewall validation.

  REDBEAR_CONFIG  Config to build (default: redbear-mini)
  --check         Automated mode — runs redbear-firewall-check inside Router-VM
                  via expect and validates all 6 scenarios pass.
  --help, -h      Show this message

Examples:
  # Build and launch 3-VM topology interactively:
  ./local/scripts/test-firewall-scenarios.sh redbear-mini

  # Run automated check (builds if needed):
  ./local/scripts/test-firewall-scenarios.sh --check redbear-mini
USAGE
}

# ──────────────────────────────────────────────────────────────────────────
# Parse arguments
# ──────────────────────────────────────────────────────────────────────────

check_mode=0
for arg in "$@"; do
    case "$arg" in
        --help|-h|help)
            usage
            exit 0
            ;;
        --check)
            check_mode=1
            ;;
        -*)
            echo "ERROR: unknown flag: $arg" >&2
            usage >&2
            exit 1
            ;;
        *)
            ISO_CONFIG="$arg"
            ;;
    esac
done

# ──────────────────────────────────────────────────────────────────────────
# Build prerequisites
# ──────────────────────────────────────────────────────────────────────────

ISO_PATH="$ROOT_DIR/build/$ARCH/$ISO_CONFIG.iso"
BUILD_SCRIPT="$ROOT_DIR/local/scripts/build-redbear.sh"

if [[ ! -f "$ISO_PATH" ]]; then
    echo "=== ISO not found, building $ISO_CONFIG ==="
    "$BUILD_SCRIPT" "$ISO_CONFIG" || {
        echo "ERROR: build failed" >&2
        exit 1
    }
fi

firmware="$(find_uefi_firmware)" || {
    echo "ERROR: no usable x86_64 UEFI firmware found" >&2
    exit 1
}

echo "UEFI firmware: $firmware"
echo "ISO: $ISO_PATH"

# ──────────────────────────────────────────────────────────────────────────
# Create artifact directory
# ──────────────────────────────────────────────────────────────────────────

mkdir -p "$ARTIFACT_DIR"
echo "Artifacts will be saved to: $ARTIFACT_DIR"

# ──────────────────────────────────────────────────────────────────────────
# Create AF_UNIX sockets for inter-VM networking
# ──────────────────────────────────────────────────────────────────────────

SOCK_A="$ARTIFACT_DIR/vm-a-net0.sock"
SOCK_B="$ARTIFACT_DIR/vm-b-net0.sock"
SOCK_RTR_INT="$ARTIFACT_DIR/router-int.sock"   # Router eth0 (internal)
SOCK_RTR_EXT="$ARTIFACT_DIR/router-ext.sock"   # Router eth1 (external)

# Cleanup any stale sockets
rm -f "$SOCK_A" "$SOCK_B" "$SOCK_RTR_INT" "$SOCK_RTR_EXT"
sock_paths=("$SOCK_A" "$SOCK_B" "$SOCK_RTR_INT" "$SOCK_RTR_EXT")

# ──────────────────────────────────────────────────────────────────────────
# Router-VM: two NICs (eth0=internal subnet, eth1=external)
# ──────────────────────────────────────────────────────────────────────────

ROUTER_ISO="$ISO_PATH"
ROUTER_PCAP_INT="$ARTIFACT_DIR/router-eth0.pcap"
ROUTER_PCAP_EXT="$ARTIFACT_DIR/router-eth1.pcap"
ROUTER_QMP="$ARTIFACT_DIR/router-qmp.sock"

# Check if expect script is available for automated mode
EXPECT_SCRIPT="$ROOT_DIR/local/scripts/qemu-login-expect.py"
if [[ "$check_mode" -eq 1 ]]; then
    if [[ ! -f "$EXPECT_SCRIPT" ]]; then
        echo "WARNING: expect script not found at $EXPECT_SCRIPT — falling back to interactive mode" >&2
        check_mode=0
    fi
    if ! command -v python3 &>/dev/null; then
        echo "WARNING: python3 not found — falling back to interactive mode" >&2
        check_mode=0
    fi
fi

echo ""
echo "=== Topology ==="
echo "  VM-A:      10.0.0.2/24 (internal subnet)"
echo "  Router-VM: 10.0.0.1/24 (eth0, internal), 10.0.2.100/24 (eth1, external)"
echo "  VM-B:      10.0.2.200/24 (external subnet)"
echo ""
echo "=== Launching VMs ==="

# ──────────────────────────────────────────────────────────────────────────
# Launch Router-VM first (firewall host, runs redbear-firewall-check)
# ──────────────────────────────────────────────────────────────────────────

echo "Starting Router-VM..."
if [[ "$check_mode" -eq 1 ]]; then
    python3 "$EXPECT_SCRIPT" \
        --timeout 300 \
        --log "$ARTIFACT_DIR/router-boot.log" \
        --step "expect:login:" \
        --step "send:root" \
        --step "expect:assword:" \
        --step "send:password" \
        --step "expect:Type 'help' for available commands." \
        --step "send:redbear-firewall-check" \
        --step "expect:Firewall Validation Harness" \
        --step "expect:SCENARIO_A=pass" \
        --step "expect:SCENARIO_B=pass" \
        --step "expect:SCENARIO_C=pass" \
        --step "expect:SCENARIO_D=pass" \
        --step "expect:SCENARIO_E=pass" \
        --step "expect:SCENARIO_F=pass" \
        --step "expect:Summary" \
        --pass_marker "SCENARIO_F=pass" \
        --step "send:shutdown" \
        -- \
        qemu-system-x86_64 \
            -name "router-vm" \
            -smp 2 \
            -m 512 \
            -bios "$firmware" \
            -machine q35 \
            -nographic \
            -cdrom "$ROUTER_ISO" \
            -netdev socket,id=net0,listen=:0,path="$SOCK_RTR_INT" \
            -device virtio-net,netdev=net0,mac=52:54:00:12:34:01 \
            -object filter-dump,id=f0,netdev=net0,file="$ROUTER_PCAP_INT" \
            -netdev socket,id=net1,listen=:0,path="$SOCK_RTR_EXT" \
            -device virtio-net,netdev=net1,mac=52:54:00:12:34:02 \
            -object filter-dump,id=f1,netdev=net1,file="$ROUTER_PCAP_EXT" \
            -qmp unix:"$ROUTER_QMP",server,nowait \
            -enable-kvm -cpu host &
    router_pid=$!
else
    qemu-system-x86_64 \
        -name "router-vm" \
        -smp 2 \
        -m 512 \
        -bios "$firmware" \
        -machine q35 \
        -nographic \
        -cdrom "$ROUTER_ISO" \
        -netdev socket,id=net0,listen=:0,path="$SOCK_RTR_INT" \
        -device virtio-net,netdev=net0,mac=52:54:00:12:34:01 \
        -object filter-dump,id=f0,netdev=net0,file="$ROUTER_PCAP_INT" \
        -netdev socket,id=net1,listen=:0,path="$SOCK_RTR_EXT" \
        -device virtio-net,netdev=net1,mac=52:54:00:12:34:02 \
        -object filter-dump,id=f1,netdev=net1,file="$ROUTER_PCAP_EXT" \
        -qmp unix:"$ROUTER_QMP",server,nowait \
        -enable-kvm -cpu host &
    router_pid=$!
fi
vm_pids+=("$router_pid")
echo "  Router-VM PID: $router_pid"
echo "  Router pcaps: $ROUTER_PCAP_INT, $ROUTER_PCAP_EXT"

# Wait a moment for router sockets to be created before connecting VMs
sleep 2
for sock in "$SOCK_RTR_INT" "$SOCK_RTR_EXT"; do
    waited=0
    while [[ ! -S "$sock" ]] && [[ $waited -lt 30 ]]; do
        sleep 1
        waited=$((waited + 1))
    done
    if [[ ! -S "$sock" ]]; then
        echo "WARNING: socket $sock not created after 30s — continuing anyway" >&2
    fi
done

# ──────────────────────────────────────────────────────────────────────────
# Launch VM-A (internal subnet peer)
# ──────────────────────────────────────────────────────────────────────────

echo "Starting VM-A (internal subnet)..."
VM_A_PCAP="$ARTIFACT_DIR/vm-a-net0.pcap"
qemu-system-x86_64 \
    -name "vm-a" \
    -smp 1 \
    -m 256 \
    -bios "$firmware" \
    -machine q35 \
    -nographic \
    -cdrom "$ISO_PATH" \
    -netdev socket,id=net0,connect=:0,path="$SOCK_RTR_INT" \
    -device virtio-net,netdev=net0,mac=52:54:00:12:34:0A \
    -object filter-dump,id=fa,netdev=net0,file="$VM_A_PCAP" \
    -enable-kvm -cpu host &
vm_a_pid=$!
vm_pids+=("$vm_a_pid")
echo "  VM-A PID: $vm_a_pid"
echo "  VM-A pcap: $VM_A_PCAP"

# ──────────────────────────────────────────────────────────────────────────
# Launch VM-B (external subnet peer)
# ──────────────────────────────────────────────────────────────────────────

echo "Starting VM-B (external subnet)..."
VM_B_PCAP="$ARTIFACT_DIR/vm-b-net0.pcap"
qemu-system-x86_64 \
    -name "vm-b" \
    -smp 1 \
    -m 256 \
    -bios "$firmware" \
    -machine q35 \
    -nographic \
    -cdrom "$ISO_PATH" \
    -netdev socket,id=net0,connect=:0,path="$SOCK_RTR_EXT" \
    -device virtio-net,netdev=net0,mac=52:54:00:12:34:0B \
    -object filter-dump,id=fb,netdev=net0,file="$VM_B_PCAP" \
    -enable-kvm -cpu host &
vm_b_pid=$!
vm_pids+=("$vm_b_pid")
echo "  VM-B PID: $vm_b_pid"
echo "  VM-B pcap: $VM_B_PCAP"

echo ""
echo "=== All 3 VMs launched ==="
echo "PIDs: router=$router_pid, vm-a=$vm_a_pid, vm-b=$vm_b_pid"
echo ""

# ──────────────────────────────────────────────────────────────────────────
# Summary file
# ──────────────────────────────────────────────────────────────────────────

SUMMARY="$ARTIFACT_DIR/summary.txt"
cat > "$SUMMARY" <<SUMMARYEOF
Firewall Validation Run
=======================
Timestamp: $(date -u +"%Y-%m-%dT%H:%M:%SZ")
Config: $ISO_CONFIG
Mode: $([ "$check_mode" -eq 1 ] && echo "automated (check)" || echo "interactive")

Topology:
  Router-VM:  MAC 52:54:00:12:34:01 (eth0, internal), MAC 52:54:00:12:34:02 (eth1, external)
  VM-A:       MAC 52:54:00:12:34:0A (internal subnet peer)
  VM-B:       MAC 52:54:00:12:34:0B (external subnet peer)

Artifacts:
  router-eth0.pcap:    Internal NIC capture
  router-eth1.pcap:    External NIC capture
  vm-a-net0.pcap:      Internal peer capture
  vm-b-net0.pcap:      External peer capture
  router-boot.log:     Router VM serial console log

PIDs:
  router-vm: $router_pid
  vm-a: $vm_a_pid
  vm-b: $vm_b_pid
SUMMARYEOF

echo "Summary written to: $SUMMARY"

if [[ "$check_mode" -eq 1 ]]; then
    echo ""
    echo "Waiting for Router-VM expect script to complete (max 300s)..."
    wait "$router_pid" || true
    router_rc=${PIPESTATUS[0]:-0}
    echo "Router-VM exited with code: $router_rc"
    echo ""

    # Check for the pass markers in the router boot log
    ROUTER_LOG="$ARTIFACT_DIR/router-boot.log"
    if [[ -f "$ROUTER_LOG" ]]; then
        all_pass=1
        for label in A B C D E F; do
            if grep -q "SCENARIO_${label}=pass" "$ROUTER_LOG" 2>/dev/null; then
                echo "SCENARIO_${label}: ✅ pass"
            else
                echo "SCENARIO_${label}: ❌ not found in output"
                all_pass=0
            fi
        done

        if [[ $all_pass -eq 1 ]]; then
            echo ""
            echo "=== All 6 firewall scenarios PASSED ==="
            exit 0
        else
            echo ""
            echo "=== Some firewall scenarios FAILED or did not produce output ==="
            echo "Full log: $ROUTER_LOG"
            exit 1
        fi
    else
        echo "ERROR: Router boot log not found at $ROUTER_LOG"
        exit 1
    fi
else
    echo ""
    echo "=== VMs are running — press Ctrl+C to stop ==="
    echo "Log into Router-VM and run: redbear-firewall-check"
    echo ""

    # Wait for any VM to exit, or for Ctrl+C
    wait -n "${vm_pids[@]}" 2>/dev/null || true
fi