#!/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

    # Allow return traffic to wg interface
    if ! iptables -C FORWARD -o "$WG_IF" -j ACCEPT 2>/dev/null; then
        iptables -I FORWARD 1 -o "$WG_IF" -j ACCEPT
        log_info "Added FORWARD rule for return traffic to $WG_IF"
    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 -A FORWARD -i "$WG_IF" -j "$CHAIN"
    fi
    log_info "Chain $CHAIN created and linked to FORWARD"

    # === 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 "$@"