#!/bin/bash
# WireGuard Policy Firewall Installer/Uninstaller

set -euo pipefail

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

install_policy() {
    echo "Installing WireGuard Policy Firewall..."

    # Install dependencies
    echo "Checking dependencies..."
    apt-get update -y || true
    apt-get install -y jq inotify-tools ipset iptables || true

    echo "Writing scripts to /usr/local/bin/..."

    cat << 'EOF_WG_POLICY_LIB' > /usr/local/bin/wg-policy-lib.sh
#!/bin/bash
# wg-policy-lib.sh — Shared functions for WireGuard Policy Firewall
# Source this file; do not execute directly.

set -euo pipefail

# ============================================================
# CONFIGURATION
# ============================================================
readonly WG_IF="${WG_IF:-wg0}"
readonly CHAIN="WG_POLICY"
readonly CHAIN_BACKUP="WG_POLICY_BAK"
readonly POLICY_FILE="/etc/wireguard/policy.json"
readonly WG_CONF="/etc/wireguard/wg0.conf"
readonly LOCK_FILE="/var/lock/wg-policy.lock"
readonly BACKUP_DIR="/etc/wireguard/backups"
readonly LOG_PREFIX="WG_DROP"
readonly LOG_RATE="10/min"
readonly IPSET_V4="wg_allowed_v4"
readonly IPSET_V6="wg_allowed_v6"
readonly MAX_RETRY=3
readonly RETRY_DELAY=2
readonly DEBOUNCE_SEC=2

# ============================================================
# LOGGING
# ============================================================
log_info()  { echo "[$(date '+%Y-%m-%d %H:%M:%S')] [INFO]  $*"; }
log_warn()  { echo "[$(date '+%Y-%m-%d %H:%M:%S')] [WARN]  $*" >&2; }
log_error() { echo "[$(date '+%Y-%m-%d %H:%M:%S')] [ERROR] $*" >&2; }

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

# Validate IPv4 address (strict: 0-255 per octet, no leading zeros)
validate_ipv4() {
    local ip="$1"
    # Match basic pattern
    if [[ ! "$ip" =~ ^([0-9]{1,3})\.([0-9]{1,3})\.([0-9]{1,3})\.([0-9]{1,3})$ ]]; then
        return 1
    fi
    local IFS='.'
    read -ra octets <<< "$ip"
    for octet in "${octets[@]}"; do
        # Reject leading zeros (except "0" itself)
        if [[ "$octet" =~ ^0[0-9] ]]; then
            return 1
        fi
        if (( octet < 0 || octet > 255 )); then
            return 1
        fi
    done
    return 0
}

# Validate IPv4 CIDR (e.g., 192.168.1.0/24)
validate_ipv4_cidr() {
    local cidr="$1"
    local ip prefix

    if [[ "$cidr" == *"/"* ]]; then
        ip="${cidr%%/*}"
        prefix="${cidr##*/}"
    else
        # Single IP treated as /32
        ip="$cidr"
        prefix="32"
    fi

    if ! validate_ipv4 "$ip"; then
        return 1
    fi

    if [[ ! "$prefix" =~ ^[0-9]+$ ]] || (( prefix < 0 || prefix > 32 )); then
        return 1
    fi
    return 0
}

# Validate IPv6 address (basic check)
validate_ipv6() {
    local ip="$1"
    # Basic IPv6 pattern — covers full, compressed, and mixed notation
    if [[ "$ip" =~ ^([0-9a-fA-F]{0,4}:){2,7}[0-9a-fA-F]{0,4}$ ]] || \
       [[ "$ip" =~ ^::([0-9a-fA-F]{0,4}:){0,5}[0-9a-fA-F]{0,4}$ ]] || \
       [[ "$ip" =~ ^([0-9a-fA-F]{0,4}:){1,7}:$ ]] || \
       [[ "$ip" == "::" ]] || \
       [[ "$ip" == "::1" ]]; then
        return 0
    fi
    return 1
}

# Validate IPv6 CIDR
validate_ipv6_cidr() {
    local cidr="$1"
    local ip prefix

    if [[ "$cidr" == *"/"* ]]; then
        ip="${cidr%%/*}"
        prefix="${cidr##*/}"
    else
        ip="$cidr"
        prefix="128"
    fi

    if ! validate_ipv6 "$ip"; then
        return 1
    fi

    if [[ ! "$prefix" =~ ^[0-9]+$ ]] || (( prefix < 0 || prefix > 128 )); then
        return 1
    fi
    return 0
}

# Generic CIDR validator — dispatches to v4 or v6
validate_cidr() {
    local cidr="$1"
    if [[ "$cidr" == *":"* ]]; then
        validate_ipv6_cidr "$cidr"
    else
        validate_ipv4_cidr "$cidr"
    fi
}

# ============================================================
# IPSET MANAGEMENT
# ============================================================

has_ipset() {
    command -v ipset &>/dev/null
}

ensure_ipset() {
    local name="$1" family="$2"
    has_ipset || return 0
    if ! ipset list "$name" &>/dev/null; then
        ipset create "$name" hash:net,net family "$family" hashsize 1024 maxelem 65536 timeout 0
        log_info "Created ipset: $name (family=$family)"
    fi
}

flush_ipset() {
    local name="$1"
    has_ipset || return 0
    if ipset list "$name" &>/dev/null; then
        ipset flush "$name"
    fi
}

destroy_ipset() {
    local name="$1"
    has_ipset || return 0
    if ipset list "$name" &>/dev/null; then
        ipset destroy "$name"
    fi
}

# ============================================================
# RETRY MECHANISM
# ============================================================

retry() {
    local max_attempts="${MAX_RETRY}"
    local delay="${RETRY_DELAY}"
    local attempt=1
    local exit_code=0

    while (( attempt <= max_attempts )); do
        if "$@"; then
            return 0
        fi
        exit_code=$?
        log_warn "Attempt $attempt/$max_attempts failed (exit=$exit_code), retrying in ${delay}s..."
        sleep "$delay"
        (( attempt++ ))
        (( delay *= 2 ))  # exponential backoff
    done

    log_error "All $max_attempts attempts failed for: $*"
    return "$exit_code"
}

# ============================================================
# LOCK MANAGEMENT
# ============================================================

acquire_lock() {
    local lock_fd=200
    eval "exec ${lock_fd}>\"${LOCK_FILE}\""
    if ! flock -x -w 10 "$lock_fd"; then
        log_error "Failed to acquire lock: ${LOCK_FILE} (timeout 10s)"
        return 1
    fi
    log_info "Lock acquired: ${LOCK_FILE}"
}

release_lock() {
    # Lock released automatically when fd closes, but we clean up file
    rm -f "$LOCK_FILE" 2>/dev/null || true
}

# ============================================================
# BACKUP
# ============================================================

backup_policy() {
    mkdir -p "$BACKUP_DIR"
    local timestamp
    timestamp="$(date '+%Y%m%d_%H%M%S')"

    if [[ -f "$POLICY_FILE" ]]; then
        cp "$POLICY_FILE" "${BACKUP_DIR}/policy_${timestamp}.json"
        log_info "Backup created: ${BACKUP_DIR}/policy_${timestamp}.json"
    fi

    # Keep only last 50 backups
    local count
    count=$(find "$BACKUP_DIR" -name 'policy_*.json' -type f | wc -l)
    if (( count > 50 )); then
        find "$BACKUP_DIR" -name 'policy_*.json' -type f -printf '%T@ %p\n' \
            | sort -n \
            | head -n $(( count - 50 )) \
            | awk '{print $2}' \
            | xargs rm -f
        log_info "Pruned old backups (kept 50)"
    fi
}

backup_iptables() {
    mkdir -p "$BACKUP_DIR"
    local timestamp
    timestamp="$(date '+%Y%m%d_%H%M%S')"

    if iptables-save > "${BACKUP_DIR}/iptables_${timestamp}.rules" 2>/dev/null; then
        log_info "iptables backup: ${BACKUP_DIR}/iptables_${timestamp}.rules"
    fi

    if command -v ip6tables-save &>/dev/null; then
        ip6tables-save > "${BACKUP_DIR}/ip6tables_${timestamp}.rules" 2>/dev/null || true
    fi

    # Keep only last 20 iptables backups
    for prefix in iptables ip6tables; do
        local count
        count=$(find "$BACKUP_DIR" -name "${prefix}_*.rules" -type f | wc -l)
        if (( count > 20 )); then
            find "$BACKUP_DIR" -name "${prefix}_*.rules" -type f -printf '%T@ %p\n' \
                | sort -n \
                | head -n $(( count - 20 )) \
                | awk '{print $2}' \
                | xargs rm -f
        fi
    done
}

# ============================================================
# DETECTION HELPERS
# ============================================================

detect_lan_subnets() {
    ip -o -4 route show scope link \
        | awk '{print $1}' \
        | grep -E '^(10\.|172\.(1[6-9]|2[0-9]|3[01])\.|192\.168\.)' \
        | sort -u
}

detect_wg_subnet() {
    local family="${1:-inet}"
    if [[ "$family" == "inet6" ]]; then
        ip -o -6 addr show "$WG_IF" 2>/dev/null \
            | awk '{print $4; exit}'
    else
        ip -o -4 addr show "$WG_IF" 2>/dev/null \
            | awk '{print $4; exit}'
    fi
}

detect_default_if() {
    local def_if
    def_if=$(ip -4 route ls 2>/dev/null | grep default | grep -Po '(?<=dev )(\S+)' | head -1 || true)
    if [[ -z "$def_if" ]]; then
        echo "eth0"
    else
        echo "$def_if"
    fi
}

# ============================================================
# HEALTH CHECK
# ============================================================

health_check() {
    local status=0
    local report=""

    # 1. Check interface exists
    if ip link show "$WG_IF" &>/dev/null; then
        report+="[OK] Interface $WG_IF is UP\n"
    else
        report+="[FAIL] Interface $WG_IF not found\n"
        status=1
    fi

    # 2. Check policy.json exists and is valid
    if [[ -f "$POLICY_FILE" ]] && jq empty "$POLICY_FILE" 2>/dev/null; then
        local client_count
        client_count=$(jq '(.clients // {}) | length' "$POLICY_FILE")
        report+="[OK] policy.json valid ($client_count clients)\n"
    else
        report+="[FAIL] policy.json missing or corrupt\n"
        status=1
    fi

    # 3. Check chain exists
    if iptables -L "$CHAIN" -n &>/dev/null; then
        local rule_count
        rule_count=$(iptables -L "$CHAIN" -n 2>/dev/null | tail -n +3 | wc -l)
        report+="[OK] Chain $CHAIN active ($rule_count rules)\n"
    else
        report+="[WARN] Chain $CHAIN not found\n"
        status=1
    fi

    # 4. Check FORWARD reference
    if iptables -L FORWARD -n 2>/dev/null | grep -q "$CHAIN"; then
        report+="[OK] FORWARD chain references $CHAIN\n"
    else
        report+="[WARN] FORWARD chain has no reference to $CHAIN\n"
        status=1
    fi

    # 5. Check ipset
    for set_name in "$IPSET_V4" "$IPSET_V6"; do
        if ipset list "$set_name" &>/dev/null; then
            local entry_count
            entry_count=$(ipset list "$set_name" 2>/dev/null | grep -c '^[0-9a-f:]' || echo 0)
            report+="[OK] ipset $set_name active ($entry_count entries)\n"
        else
            report+="[INFO] ipset $set_name not created (may not be needed)\n"
        fi
    done

    # 6. Check watcher service
    if systemctl is-active --quiet wg-policy.service 2>/dev/null; then
        report+="[OK] wg-policy.service is running\n"
    else
        report+="[INFO] wg-policy.service not running\n"
    fi

    # 7. Check lock file not stale
    if [[ -f "$LOCK_FILE" ]]; then
        local lock_age
        lock_age=$(( $(date +%s) - $(stat -c %Y "$LOCK_FILE" 2>/dev/null || echo 0) ))
        if (( lock_age > 300 )); then
            report+="[WARN] Stale lock file (${lock_age}s old)\n"
        else
            report+="[OK] Lock file age: ${lock_age}s\n"
        fi
    else
        report+="[OK] No stale lock file\n"
    fi

    echo -e "$report"
    return $status
}
EOF_WG_POLICY_LIB

    cat << 'EOF_WG_POLICY_ENGINE' > /usr/local/bin/wg-policy-engine.sh
#!/bin/bash
# wg-policy-engine.sh — Applies iptables/ipset rules from policy.json
# Fixed: unquoted variables, LOG placement, LAN block targeting,
#        atomic chain swap, ipset, IPv6 optional, rollback on failure

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "${SCRIPT_DIR}/wg-policy-lib.sh"

# ============================================================
# ROLLBACK
# ============================================================

rollback() {
    log_error "ROLLBACK triggered! Restoring previous rules..."

    # Remove new chain references
    while true; do
        local rline=""
        rline=$(iptables -nL FORWARD --line-numbers 2>/dev/null | grep "$CHAIN" | awk '{print $1}' | head -n 1 || true)
        if [[ -n "$rline" ]]; then
            iptables -D FORWARD "$rline" 2>/dev/null || break
        else
            break
        fi
    done

    # Flush and remove new chain
    iptables -F "$CHAIN" 2>/dev/null || true
    iptables -X "$CHAIN" 2>/dev/null || true

    # Restore backup chain if it exists
    if iptables -L "$CHAIN_BACKUP" -n &>/dev/null; then
        # Rename backup chain to active
        iptables -N "$CHAIN" 2>/dev/null || iptables -F "$CHAIN"
        # Copy rules from backup
        iptables-save -c | grep "^-A $CHAIN_BACKUP" | \
            sed "s/-A $CHAIN_BACKUP/-A $CHAIN/" | \
            iptables-restore -c 2>/dev/null || true

        iptables -A FORWARD -i "$WG_IF" -j "$CHAIN"
        log_info "Rollback: restored from backup chain"
    fi

    # Cleanup backup chain
    iptables -F "$CHAIN_BACKUP" 2>/dev/null || true
    iptables -X "$CHAIN_BACKUP" 2>/dev/null || true

    # Cleanup backup ipsets
    destroy_ipset "${IPSET_V4}_bak" 2>/dev/null || true
    destroy_ipset "${IPSET_V6}_bak" 2>/dev/null || true
}

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

main() {
    log_info "Starting policy engine..."

    # === VALIDATE ===
    if ! ip link show "$WG_IF" &>/dev/null; then
        log_error "Interface $WG_IF is not running. Aborting policy engine."
        exit 1
    fi

    if [[ ! -f "$POLICY_FILE" ]]; then
        log_error "Policy file not found: $POLICY_FILE"
        exit 1
    fi

    if ! jq empty "$POLICY_FILE" 2>/dev/null; then
        log_error "policy.json is corrupt"
        exit 1
    fi

    # Backup iptables state
    backup_iptables

    # Set trap for rollback on failure
    trap 'rollback' ERR

    # === DETECT SUBNETS ===
    local WG_SUBNET WG_SUBNET_V6 LAN_SUBNETS DEF_IF

    WG_SUBNET="$(detect_wg_subnet inet)"
    WG_SUBNET_V6="$(detect_wg_subnet inet6)"
    LAN_SUBNETS="$(detect_lan_subnets)"
    DEF_IF="$(detect_default_if)"

    if [[ -z "$WG_SUBNET" ]]; then
        log_warn "Interface $WG_IF has no IPv4, skipping client isolation"
    else
        log_info "WG IPv4 subnet: $WG_SUBNET"
    fi

    if [[ -n "$WG_SUBNET_V6" ]]; then
        log_info "WG IPv6 subnet: $WG_SUBNET_V6"
    fi

    if [[ -n "$LAN_SUBNETS" ]]; then
        log_info "Detected LAN subnets:"
        echo "$LAN_SUBNETS" | while read -r s; do log_info "  $s"; done
    fi

    # === BASE ROUTING & NAT ===
    # Enable IP Forwarding
    sysctl -w net.ipv4.ip_forward=1 >/dev/null 2>&1 || true
    if command -v ip6tables &>/dev/null; then
        sysctl -w net.ipv6.conf.all.forwarding=1 >/dev/null 2>&1 || true
    fi

    # Setup MASQUERADE on default interface
    if ! iptables -t nat -C POSTROUTING -o "$DEF_IF" -j MASQUERADE 2>/dev/null; then
        iptables -t nat -A POSTROUTING -o "$DEF_IF" -j MASQUERADE
        log_info "Enabled IPv4 MASQUERADE on $DEF_IF"
    fi

    # === CLEANUP OLD CHAIN (loop until all references removed) ===
    log_info "Cleaning up old chain references..."
    while true; do
        local rline=""
        rline=$(iptables -nL FORWARD --line-numbers 2>/dev/null | grep "$CHAIN" | awk '{print $1}' | head -n 1 || true)
        if [[ -n "$rline" ]]; then
            iptables -D FORWARD "$rline" 2>/dev/null || break
        else
            break
        fi
    done

    # Backup existing chain before flushing
    if iptables -L "$CHAIN" -n &>/dev/null; then
        iptables -N "$CHAIN_BACKUP" 2>/dev/null || iptables -F "$CHAIN_BACKUP"
        iptables-save -c 2>/dev/null | grep "^-A $CHAIN" | \
            sed "s/-A $CHAIN/-A $CHAIN_BACKUP/" | \
            iptables-restore -c 2>/dev/null || true
        log_info "Backed up existing chain to $CHAIN_BACKUP"
    fi

    iptables -F "$CHAIN" 2>/dev/null || true
    iptables -X "$CHAIN" 2>/dev/null || true

    # === CREATE FRESH CHAIN ===
    iptables -N "$CHAIN"

    if ! iptables -C FORWARD -i "$WG_IF" -j "$CHAIN" 2>/dev/null; then
        iptables -I FORWARD 1 -i "$WG_IF" -j "$CHAIN"
    fi
    
    if ! iptables -C FORWARD -o "$WG_IF" -j "$CHAIN" 2>/dev/null; then
        iptables -I FORWARD 2 -o "$WG_IF" -j "$CHAIN"
    fi
    log_info "Chain $CHAIN created and linked to FORWARD (In/Out)"

    # === POPULATE IPSET (hash:net,net for source->target mapping) ===
    local use_ipset=false
    if has_ipset; then
        use_ipset=true
        log_info "Populating ipsets..."

        ensure_ipset "$IPSET_V4" "inet"
        flush_ipset "$IPSET_V4"

        # Check if we need IPv6 ipset
        local use_ipv6=false
        if [[ -n "$WG_SUBNET_V6" ]] && command -v ip6tables &>/dev/null; then
            use_ipv6=true
            ensure_ipset "$IPSET_V6" "inet6"
            flush_ipset "$IPSET_V6"
        fi

        # Read all access entries and populate ipset (client_ip,target)
        jq -r '
            .clients // {} | to_entries[] |
            select(.value.access != null and (.value.access | length > 0)) |
            .key as $ip |
            .value.access[] |
            "\($ip) \(.)"
        ' "$POLICY_FILE" 2>/dev/null | while read -r client_ip target; do
            [[ -z "$client_ip" || -z "$target" ]] && continue

            if [[ "$target" == *":"* ]]; then
                if [[ "$use_ipv6" == true ]]; then
                    ipset add "$IPSET_V6" "${client_ip},${target}" 2>/dev/null || \
                        log_warn "Failed to add ${client_ip},${target} to ipset $IPSET_V6"
                fi
            else
                ipset add "$IPSET_V4" "${client_ip},${target}" 2>/dev/null || \
                    log_warn "Failed to add ${client_ip},${target} to ipset $IPSET_V4"
            fi
        done

        local v4_count v6_count
        v4_count=$(ipset list "$IPSET_V4" 2>/dev/null | grep -c '^[0-9]' || echo 0)
        v6_count=$(ipset list "$IPSET_V6" 2>/dev/null | grep -c '^[0-9a-f:]' || echo 0)
        log_info "ipset $IPSET_V4: $v4_count entries, $IPSET_V6: $v6_count entries"
    else
        log_warn "ipset not installed, falling back to per-rule iptables whitelist"
        local use_ipv6=false
        if [[ -n "$WG_SUBNET_V6" ]] && command -v ip6tables &>/dev/null; then
            use_ipv6=true
        fi
    fi

    # === RULE 1: ESTABLISHED,RELATED — allow return traffic ===
    iptables -A "$CHAIN" -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT

    # === RULE 2: WHITELIST (per-client source) ===
    if [[ "$use_ipset" == true ]]; then
        iptables -A "$CHAIN" -m set --match-set "$IPSET_V4" src,dst -j ACCEPT
        if [[ "$use_ipv6" == true ]]; then
            ip6tables -A "$CHAIN" -m set --match-set "$IPSET_V6" src,dst -j ACCEPT 2>/dev/null || true
        fi
    else
        jq -r '
            .clients // {} | to_entries[] |
            select(.value.access != null and (.value.access | length > 0)) |
            .key as $ip |
            .value.access[] |
            "\($ip) \(.)"
        ' "$POLICY_FILE" 2>/dev/null | while read -r client_ip target; do
            [[ -z "$client_ip" || -z "$target" ]] && continue

            if [[ "$client_ip" == *":"* ]]; then
                if [[ "$use_ipv6" == true ]]; then
                    ip6tables -A "$CHAIN" -s "$client_ip" -d "$target" -j ACCEPT 2>/dev/null || true
                fi
            else
                iptables -A "$CHAIN" -s "$client_ip" -d "$target" -j ACCEPT
            fi
        done
    fi

    # === RULE 3: ISOLATION — drop NEW connections between WG clients ===
    if [[ -n "$WG_SUBNET" ]]; then
        iptables -A "$CHAIN" \
            -s "$WG_SUBNET" \
            -d "$WG_SUBNET" \
            -m conntrack --ctstate NEW \
            -j DROP
        log_info "Client isolation enabled for $WG_SUBNET"
    fi

    if [[ "$use_ipv6" == true && -n "$WG_SUBNET_V6" ]]; then
        ip6tables -A "$CHAIN" \
            -s "$WG_SUBNET_V6" \
            -d "$WG_SUBNET_V6" \
            -m conntrack --ctstate NEW \
            -j DROP 2>/dev/null || true
        log_info "Client isolation enabled for IPv6 $WG_SUBNET_V6"
    fi

    # === RULE 4: BLOCK LAN — drop from WG subnet to private LAN ===
    if [[ -n "$WG_SUBNET" && -n "$LAN_SUBNETS" ]]; then
        echo "$LAN_SUBNETS" | while read -r subnet; do
            [[ -z "$subnet" ]] && continue
            # Skip if LAN subnet exactly matches WG subnet (handled by Rule 3)
            [[ "$subnet" == "$WG_SUBNET" ]] && continue

            iptables -A "$CHAIN" -s "$WG_SUBNET" -d "$subnet" -j DROP
            log_info "Block: $WG_SUBNET → $subnet"
        done
    fi

    # IPv6 LAN block (link-local and ULA)
    if [[ "$use_ipv6" == true && -n "$WG_SUBNET_V6" ]]; then
        # Block to link-local (fe80::/10)
        ip6tables -A "$CHAIN" -s "$WG_SUBNET_V6" -d "fe80::/10" -j DROP 2>/dev/null || true
        # Block to ULA (fc00::/7)
        ip6tables -A "$CHAIN" -s "$WG_SUBNET_V6" -d "fc00::/7" -j DROP 2>/dev/null || true
        log_info "IPv6 LAN block applied (link-local + ULA)"
    fi

    # === RULE 5: INTERNET ACCESS (#Internet = true) ===
    jq -r '
        .clients // {} | to_entries[] |
        select(.value.internet == "true") |
        "\(.key)"
    ' "$POLICY_FILE" 2>/dev/null | while read -r client_ip; do
        [[ -z "$client_ip" ]] && continue

        if [[ "$client_ip" == *":"* ]]; then
            if [[ "$use_ipv6" == true ]]; then
                ip6tables -A "$CHAIN" -s "$client_ip" -j ACCEPT 2>/dev/null || true
            fi
        else
            iptables -A "$CHAIN" -s "$client_ip" -j ACCEPT
        fi
    done

    # === RULE 6: LOGGING (rate-limited) — BEFORE final DROP ===
    iptables -A "$CHAIN" \
        -m limit --limit "$LOG_RATE" \
        -j LOG --log-prefix "${LOG_PREFIX}: " --log-level 4

    if [[ "$use_ipv6" == true ]]; then
        ip6tables -A "$CHAIN" \
            -m limit --limit "$LOG_RATE" \
            -j LOG --log-prefix "${LOG_PREFIX}: " --log-level 4 2>/dev/null || true
    fi

    # === RULE 7: DEFAULT DROP (internet block by default) ===
    iptables -A "$CHAIN" -j DROP

    if [[ "$use_ipv6" == true ]]; then
        ip6tables -A "$CHAIN" -j DROP 2>/dev/null || true
    fi

    # === CLEANUP BACKUP CHAIN (no rollback needed anymore) ===
    iptables -F "$CHAIN_BACKUP" 2>/dev/null || true
    iptables -X "$CHAIN_BACKUP" 2>/dev/null || true

    # Disable ERR trap (success path)
    trap - ERR

    # === VERIFY ===
    local rule_count
    rule_count=$(iptables -L "$CHAIN" -n 2>/dev/null | tail -n +3 | wc -l)
    log_info "Policy applied. Chain: $CHAIN, Rules: $rule_count"

    echo "[OK] iptables policy applied. Chain: $CHAIN"
}

main "$@"
EOF_WG_POLICY_ENGINE

    cat << 'EOF_WG_POLICY_CLEANUP' > /usr/local/bin/wg-policy-cleanup.sh
#!/bin/bash
# wg-policy-cleanup.sh — Clean removal of all policy artifacts

set -uo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "${SCRIPT_DIR}/wg-policy-lib.sh"

main() {
    log_info "Starting cleanup..."

    local DEF_IF
    DEF_IF="$(detect_default_if)"

    # === Base Routing Cleanup ===
    while iptables -D FORWARD -o "$WG_IF" -j ACCEPT 2>/dev/null; do :; done
    while iptables -t nat -D POSTROUTING -o "$DEF_IF" -j MASQUERADE 2>/dev/null; do :; done
    log_info "Removed base routing and NAT rules"

    # === IPv4 chain cleanup ===
    local removed=0

    while true; do
        local line=""
        line=$(iptables -nL FORWARD --line-numbers 2>/dev/null | grep "$CHAIN" | awk '{print $1}' | head -n 1 || true)
        if [[ -n "$line" ]]; then
            iptables -D FORWARD "$line" 2>/dev/null || break
            (( removed++ ))
        else
            break
        fi
    done

    if (( removed > 0 )); then
        log_info "Removed $removed FORWARD references"
    fi

    iptables -F "$CHAIN" 2>/dev/null || true
    iptables -X "$CHAIN" 2>/dev/null || true

    # Cleanup backup chain too
    iptables -F "$CHAIN_BACKUP" 2>/dev/null || true
    iptables -X "$CHAIN_BACKUP" 2>/dev/null || true

    # === IPv6 chain cleanup ===
    if command -v ip6tables &>/dev/null; then
        while true; do
            local line6=""
            line6=$(ip6tables -nL FORWARD --line-numbers 2>/dev/null | grep "$CHAIN" | awk '{print $1}' | head -n 1 || true)
            if [[ -n "$line6" ]]; then
                ip6tables -D FORWARD "$line6" 2>/dev/null || break
            else
                break
            fi
        done
        ip6tables -F "$CHAIN" 2>/dev/null || true
        ip6tables -X "$CHAIN" 2>/dev/null || true
        ip6tables -F "$CHAIN_BACKUP" 2>/dev/null || true
        ip6tables -X "$CHAIN_BACKUP" 2>/dev/null || true
    fi

    # === ipset cleanup ===
    destroy_ipset "$IPSET_V4" 2>/dev/null || true
    destroy_ipset "$IPSET_V6" 2>/dev/null || true
    destroy_ipset "${IPSET_V4}_bak" 2>/dev/null || true
    destroy_ipset "${IPSET_V6}_bak" 2>/dev/null || true

    # === Lock cleanup ===
    rm -f "$LOCK_FILE" 2>/dev/null || true

    log_info "Cleanup complete"
}

main "$@"
EOF_WG_POLICY_CLEANUP

    cat << 'EOF_WG_SYNC_POLICY' > /usr/local/bin/wg-sync-policy.sh
#!/bin/bash
# wg-sync-policy.sh — Reads wg0.conf, validates, generates policy.json atomically
# Fixed: IP validation, atomic write, proper locking, error handling

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "${SCRIPT_DIR}/wg-policy-lib.sh"

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

main() {
    log_info "Starting policy sync..."

    # Validate prerequisites
    if [[ ! -f "$WG_CONF" ]]; then
        log_error "WireGuard config not found: $WG_CONF"
        exit 1
    fi

    if ! command -v jq &>/dev/null; then
        log_error "jq is required but not installed"
        exit 1
    fi

    # Acquire lock
    acquire_lock
    trap 'release_lock' EXIT

    # Backup current policy
    backup_policy

    # Temporary file for atomic write
    local tmp_policy
    tmp_policy="$(mktemp /tmp/wg-policy.XXXXXX)"
    trap "rm -f \"$tmp_policy\" 2>/dev/null; release_lock" EXIT

    echo '{"clients":{}}' > "$tmp_policy"

    # Parse peers from wg0.conf
    # AWK extracts IP and #Access comment per [Peer] block
    local parse_errors=0

    awk '
    BEGIN { RS="\n\\[Peer\\]\n"; FS="\n" }
    NR>1 {
        ip=""; access=""; internet="false"
        for(i=1;i<=NF;i++){
            if($i ~ /^AllowedIPs/) {
                split($i,a," = ")
                gsub(/ /,"",a[2])
                split(a[2],b,",")
                split(b[1],c,"/")
                ip=c[1]
            }
            if($i ~ /^#Access/) {
                sub(/^#Access[ \t]*=?[ \t]*/, "", $i)
                access=$i
            }
            if($i ~ /^#Internet/) {
                if(tolower($i) ~ /true|yes|1|allow/) {
                    internet="true"
                }
            }
        }
        if(ip!="" && ip!="0.0.0.0" && ip!="::") {
            printf "%s|%s|%s\n", ip, access, internet
        }
    }
    ' "$WG_CONF" | while IFS="|" read -r ip access_string internet_flag; do

        # === VALIDATE CLIENT IP ===
        if ! validate_cidr "$ip"; then
            log_warn "Invalid client IP skipped: '$ip'"
            (( parse_errors++ )) || true
            continue
        fi

        # === PARSE AND VALIDATE ACCESS TARGETS ===
        local ACCESS_JSON="[]"

        if [[ -n "$access_string" ]]; then
            # Split by ; and , then validate each entry
            local valid_targets=()
            local IFS_OLD="$IFS"
            IFS=';,'
            read -ra targets <<< "$access_string"
            IFS="$IFS_OLD"

            for target in "${targets[@]}"; do
                # Trim whitespace
                target="$(echo "$target" | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')"

                [[ -z "$target" ]] && continue

                if validate_cidr "$target"; then
                    valid_targets+=("$target")
                else
                    log_warn "Invalid access target skipped for $ip: '$target'"
                    (( parse_errors++ )) || true
                fi
            done

            if (( ${#valid_targets[@]} > 0 )); then
                ACCESS_JSON=$(printf '%s\n' "${valid_targets[@]}" | jq -R . | jq -s .)
            fi
        fi

        # Write to temp policy
        jq --arg ip "$ip" --argjson access "$ACCESS_JSON" --argjson internet "$internet_flag" \
           '.clients[$ip] = {"name": $ip, "access": $access, "internet": $internet}' \
           "$tmp_policy" > "${tmp_policy}.tmp" && mv "${tmp_policy}.tmp" "$tmp_policy"

    done

    # Validate JSON before atomic move
    if ! jq empty "$tmp_policy" 2>/dev/null; then
        log_error "Generated JSON is invalid, aborting. Check $tmp_policy"
        exit 1
    fi

    # Atomic move (same filesystem = atomic rename)
    mv -f "$tmp_policy" "$POLICY_FILE"
    log_info "policy.json updated successfully"

    if (( parse_errors > 0 )); then
        log_warn "$parse_errors validation errors encountered (see warnings above)"
    fi

    local client_count
    client_count=$(jq '(.clients // {}) | length' "$POLICY_FILE")
    log_info "Total clients in policy: $client_count"
}

main "$@"
EOF_WG_SYNC_POLICY

    cat << 'EOF_WG_SYNC_WATCH' > /usr/local/bin/wg-sync-watch.sh
#!/bin/bash
# wg-sync-watch.sh — Watches wg0.conf for changes with debounce
# Fixed: proper debounce, error isolation, health reporting

set -uo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "${SCRIPT_DIR}/wg-policy-lib.sh"

SYNC_SCRIPT="/usr/local/bin/wg-sync-policy.sh"
ENGINE_SCRIPT="/usr/local/bin/wg-policy-engine.sh"
HEALTH_INTERVAL=300  # Health check every 5 minutes
LAST_HEALTH=0

# === Validate prerequisites ===
if ! command -v inotifywait &>/dev/null; then
    log_error "inotifywait not found. Install: apt install inotify-tools"
    exit 1
fi

if [[ ! -f "$WG_CONF" ]]; then
    log_error "WireGuard config not found: $WG_CONF"
    exit 1
fi

# === Main watcher loop ===
log_info "Monitoring $WG_CONF for changes (debounce: ${DEBOUNCE_SEC}s)..."
log_info "Health check interval: ${HEALTH_INTERVAL}s"

inotifywait -m -e close_write,move,create \
    --format '%e %f' \
    "$(dirname "$WG_CONF")" 2>/dev/null | \
while read -r events filename; do

    # Only react to wg0.conf changes
    [[ "$filename" != "$(basename "$WG_CONF")" ]] && continue

    log_info "Detected change: $events $filename"

    # Debounce: wait until no more events for DEBOUNCE_SEC
    while IFS= read -r -t "$DEBOUNCE_SEC" _dummy; do
        : # Drain events within debounce window
    done

    log_info "Debounce complete, applying changes..."

    # Run sync
    if retry "$SYNC_SCRIPT"; then
        log_info "Sync successful, running engine..."

        # Run engine with retry
        if retry "$ENGINE_SCRIPT"; then
            log_info "Policy engine applied successfully"
        else
            log_error "Policy engine FAILED after retries"
        fi
    else
        log_error "Policy sync FAILED after retries"
    fi

    # Periodic health check
    local now
    now=$(date +%s)
    if (( now - LAST_HEALTH >= HEALTH_INTERVAL )); then
        LAST_HEALTH=$now
        log_info "=== Periodic Health Check ==="
        health_check || log_warn "Health check reported issues"
    fi

done
EOF_WG_SYNC_WATCH

    cat << 'EOF_WG_POLICY_CTL' > /usr/local/bin/wg-policy-ctl
#!/bin/bash
# wg-policy-ctl — CLI management tool for WireGuard Policy Firewall

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "${SCRIPT_DIR}/wg-policy-lib.sh"

usage() {
    cat <<EOF
Usage: $(basename "$0") <command>

Commands:
    status      Show full health check report
    policy      Display current policy.json formatted
    rules       Show current iptables rules in WG_POLICY chain
    ipset       Show ipset contents
    log         Tail WG_DROP logs (last 50 lines)
    reload      Force re-sync and re-apply policy
    backup      Manual backup of policy + iptables
    stats       Show connection and rule statistics
    validate    Validate policy.json without applying
    help        Show this help
EOF
}

cmd_status() {
    echo "========================================="
    echo " WireGuard Policy Firewall Status"
    echo " $(date '+%Y-%m-%d %H:%M:%S')"
    echo "========================================="
    echo ""
    health_check
}

cmd_policy() {
    if [[ -f "$POLICY_FILE" ]]; then
        jq '.' "$POLICY_FILE"
    else
        log_error "policy.json not found"
        exit 1
    fi
}

cmd_rules() {
    echo "=== IPv4 Chain: $CHAIN ==="
    if iptables -L "$CHAIN" -n -v --line-numbers 2>/dev/null; then
        echo ""
    else
        echo "(chain not found)"
    fi

    echo "=== FORWARD references ==="
    iptables -L FORWARD -n -v --line-numbers 2>/dev/null | grep -i "$CHAIN" || echo "(none)"

    if command -v ip6tables &>/dev/null; then
        echo ""
        echo "=== IPv6 Chain: $CHAIN ==="
        ip6tables -L "$CHAIN" -n -v --line-numbers 2>/dev/null || echo "(chain not found)"
    fi
}

cmd_ipset() {
    for set_name in "$IPSET_V4" "$IPSET_V6"; do
        echo "=== ipset: $set_name ==="
        if ipset list "$set_name" 2>/dev/null; then
            echo ""
        else
            echo "(not found)"
            echo ""
        fi
    done
}

cmd_log() {
    echo "=== Recent WG_DROP log entries ==="
    (journalctl -k --no-pager -n 50 2>/dev/null || dmesg | tail -50) | grep "$LOG_PREFIX" || echo "(no entries)"
}

cmd_reload() {
    log_info "Force reloading policy..."
    if retry /usr/local/bin/wg-sync-policy.sh; then
        if retry /usr/local/bin/wg-policy-engine.sh; then
            log_info "Reload complete"
        else
            log_error "Engine failed"
            exit 1
        fi
    else
        log_error "Sync failed"
        exit 1
    fi
}

cmd_backup() {
    backup_policy
    backup_iptables
    log_info "Manual backup complete. Files in: $BACKUP_DIR"
}

cmd_stats() {
    echo "=== Client Count ==="
    jq '(.clients // {}) | length' "$POLICY_FILE" 2>/dev/null || echo "N/A"

    echo ""
    echo "=== Clients with Access ==="
    jq -r '.clients // {} | to_entries[] | select(.value.access | length > 0) | "\(.key): \(.value.access | join(", "))"' "$POLICY_FILE" 2>/dev/null || echo "N/A"

    echo ""
    echo "=== Clients without Access (Internet Only) ==="
    jq -r '.clients // {} | to_entries[] | select(.value.access | length == 0) | .key' "$POLICY_FILE" 2>/dev/null || echo "N/A"

    echo ""
    echo "=== Active iptables rules ==="
    iptables -L "$CHAIN" -n 2>/dev/null | tail -n +3 | wc -l || echo "N/A"

    echo ""
    echo "=== Drop count (since boot) ==="
    iptables -L "$CHAIN" -n -v 2>/dev/null | grep "DROP" | awk '{sum += $1} END {print sum+0, "packets dropped"}'

    echo ""
    echo "=== ipset entries ==="
    for set_name in "$IPSET_V4" "$IPSET_V6"; do
        local count
        count=$(ipset list "$set_name" 2>/dev/null | grep -c '^[0-9a-f:]' || echo 0)
        echo "  $set_name: $count entries"
    done
}

cmd_validate() {
    log_info "Validating policy.json..."

    if [[ ! -f "$POLICY_FILE" ]]; then
        log_error "File not found: $POLICY_FILE"
        exit 1
    fi

    if ! jq empty "$POLICY_FILE" 2>/dev/null; then
        log_error "Invalid JSON"
        exit 1
    fi

    local errors=0
    local total=0

    jq -r '.clients // {} | to_entries[] | "\(.key)|\(.value.access // [] | join(","))"' "$POLICY_FILE" | \
    while IFS="|" read -r ip access_str; do
        (( total++ ))

        if ! validate_cidr "$ip"; then
            log_error "Invalid client IP: $ip"
            (( errors++ )) || true
        fi

        if [[ -n "$access_str" ]]; then
            IFS=',' read -ra targets <<< "$access_str"
            for target in "${targets[@]}"; do
                if ! validate_cidr "$target"; then
                    log_error "Invalid access target for $ip: $target"
                    (( errors++ )) || true
                fi
            done
        fi
    done

    if (( errors > 0 )); then
        log_error "Validation failed: $errors errors"
        exit 1
    else
        log_info "Validation passed: $total clients, 0 errors"
    fi
}

# === DISPATCH ===
case "${1:-help}" in
    status)   cmd_status ;;
    policy)   cmd_policy ;;
    rules)    cmd_rules ;;
    ipset)    cmd_ipset ;;
    log)      cmd_log ;;
    reload)   cmd_reload ;;
    backup)   cmd_backup ;;
    stats)    cmd_stats ;;
    validate) cmd_validate ;;
    help|*)   usage ;;
esac
EOF_WG_POLICY_CTL

    # Write systemd files
    echo "Writing systemd units to /etc/systemd/system/..."

    cat << 'EOF_WG_POLICY_SERVICE' > /etc/systemd/system/wg-policy.service
[Unit]
Description=WireGuard Dynamic Policy Firewall Watcher
After=network-online.target wg-quick@wg0.service
Wants=wg-quick@wg0.service network-online.target
StartLimitIntervalSec=60
StartLimitBurst=5

[Service]
Type=simple
ExecStartPre=/usr/local/bin/wg-sync-policy.sh
ExecStart=/usr/local/bin/wg-sync-watch.sh
ExecStopPost=/usr/local/bin/wg-policy-cleanup.sh
Restart=always
RestartSec=10
User=root
StandardOutput=journal
StandardError=journal
SyslogIdentifier=wg-policy

# Hardening
ProtectSystem=strict
ReadWritePaths=/etc/wireguard /var/lock /tmp
ProtectHome=yes
NoNewPrivileges=no
PrivateTmp=yes

[Install]
WantedBy=multi-user.target
EOF_WG_POLICY_SERVICE

    cat << 'EOF_WG_POLICY_HEALTH_SERVICE' > /etc/systemd/system/wg-policy-health.service
[Unit]
Description=WireGuard Policy Health Check

[Service]
Type=oneshot
ExecStart=/bin/bash -c 'source /usr/local/bin/wg-policy-lib.sh && health_check'
StandardOutput=journal
StandardError=journal
SyslogIdentifier=wg-policy-health
EOF_WG_POLICY_HEALTH_SERVICE

    cat << 'EOF_WG_POLICY_HEALTH_TIMER' > /etc/systemd/system/wg-policy-health.timer
[Unit]
Description=WireGuard Policy Health Check Timer

[Timer]
OnBootSec=60
OnUnitActiveSec=300
AccuracySec=30

[Install]
WantedBy=timers.target
EOF_WG_POLICY_HEALTH_TIMER

    # Make executable
    chmod +x /usr/local/bin/wg-*.sh /usr/local/bin/wg-policy-ctl

    # Start services
    echo "Reloading systemd daemon..."
    systemctl daemon-reload
    
    echo "Enabling and starting services..."
    systemctl enable --now wg-policy.service
    systemctl enable --now wg-policy-health.timer
    
    echo "Installation complete!"
    echo "You can check status with: wg-policy-ctl status"
}

uninstall_policy() {
    echo "Uninstalling WireGuard Policy Firewall..."

    echo "Stopping and disabling services..."
    systemctl disable --now wg-policy.service wg-policy-health.timer wg-policy-health.service 2>/dev/null || true

    echo "Running cleanup script..."
    if [ -x /usr/local/bin/wg-policy-cleanup.sh ]; then
        /usr/local/bin/wg-policy-cleanup.sh || true
    fi

    echo "Removing systemd units..."
    rm -f /etc/systemd/system/wg-policy.service
    rm -f /etc/systemd/system/wg-policy-health.service
    rm -f /etc/systemd/system/wg-policy-health.timer
    systemctl daemon-reload

    echo "Removing scripts from /usr/local/bin/..."
    rm -f /usr/local/bin/wg-policy-lib.sh
    rm -f /usr/local/bin/wg-sync-policy.sh
    rm -f /usr/local/bin/wg-policy-engine.sh
    rm -f /usr/local/bin/wg-policy-cleanup.sh
    rm -f /usr/local/bin/wg-sync-watch.sh
    rm -f /usr/local/bin/wg-policy-ctl

    echo "Uninstallation complete!"
}

case "${1:-}" in
    install)
        install_policy
        ;;
    uninstall)
        uninstall_policy
        ;;
    *)
        echo "Usage: $0 {install|uninstall}"
        exit 1
        ;;
esac
