#!/bin/bash
# wg-policy-engine.sh — Applies nftables rules from policy.json

set -euo pipefail

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

NFT_FILE="/tmp/wg-policy.nft"
NFT_BACKUP="/tmp/wg-policy-backup.nft"

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

rollback() {
    log_error "ROLLBACK triggered! Restoring previous ruleset..."
    trap - ERR

    nft delete table "$NFT_TABLE_FULL" 2>/dev/null || true

    if [[ -f "$NFT_BACKUP" ]]; then
        if nft -f "$NFT_BACKUP" 2>/dev/null; then
            log_info "Rollback: restored from backup"
        else
            log_error "Rollback: failed to restore from backup"
        fi
    fi
}

# ============================================================
# RULESET GENERATION
# ============================================================

generate_ruleset() {
    local WG_SUBNET="$1"
    local WG_SUBNET_V6="$2"
    local LAN_SUBNETS="$3"
    local DEF_IF="$4"

    cat > "$NFT_FILE" << 'HEADER'
#!/usr/sbin/nft -f
flush ruleset

table inet wg_policy {

    set wg_allowed_v4 {
        type ipv4_addr . ipv4_addr
        flags interval
    }

    set wg_allowed_v6 {
        type ipv6_addr . ipv6_addr
        flags interval
    }

    chain forward {
        type filter hook forward priority filter; policy accept;

        ct state established,related accept

        iifname "wg0" ip saddr . ip daddr @wg_allowed_v4 accept
        iifname "wg0" ip6 saddr . ip6 daddr @wg_allowed_v6 accept

HEADER

    # Client isolation (IPv4)
    if [[ -n "$WG_SUBNET" ]]; then
        cat >> "$NFT_FILE" << EOF

        /* isolation: WG client to WG client */
        iifname "wg0" ip saddr $WG_SUBNET ip daddr $WG_SUBNET ct state new drop
EOF
    fi

    # Client isolation (IPv6)
    if [[ -n "$WG_SUBNET_V6" ]]; then
        cat >> "$NFT_FILE" << EOF

        /* IPv6 isolation */
        iifname "wg0" ip6 saddr $WG_SUBNET_V6 ip6 daddr $WG_SUBNET_V6 ct state new drop
        iifname "wg0" ip6 saddr $WG_SUBNET_V6 ip6 daddr fe80::/10 drop
        iifname "wg0" ip6 saddr $WG_SUBNET_V6 ip6 daddr fc00::/7 drop
EOF
    fi

    # LAN block (IPv4)
    if [[ -n "$WG_SUBNET" && -n "$LAN_SUBNETS" ]]; then
        local lan_list=""
        while IFS= read -r subnet; do
            [[ -z "$subnet" ]] && continue
            [[ "$subnet" == "$WG_SUBNET" ]] && continue
            if [[ -n "$lan_list" ]]; then
                lan_list+=", $subnet"
            else
                lan_list="$subnet"
            fi
        done <<< "$LAN_SUBNETS"

        if [[ -n "$lan_list" ]]; then
            cat >> "$NFT_FILE" << EOF

        /* LAN block */
        iifname "wg0" ip saddr $WG_SUBNET ip daddr { $lan_list } drop
EOF
        fi
    fi

    # Internet access per client (inline in chain)
    while IFS= read -r client_ip; do
        [[ -z "$client_ip" ]] && continue
        cat >> "$NFT_FILE" << EOF

        /* internet access: $client_ip */
        iifname "wg0" ip saddr $client_ip accept
EOF
    done < <(jq -r '
        .clients // {} | to_entries[] |
        select(.value.internet == "true") |
        .key
    ' "$POLICY_FILE" 2>/dev/null)

    while IFS= read -r client_ip; do
        [[ -z "$client_ip" ]] && continue
        cat >> "$NFT_FILE" << EOF

        /* internet access v6: $client_ip */
        iifname "wg0" ip6 saddr $client_ip accept
EOF
    done < <(jq -r '
        .clients // {} | to_entries[] |
        select(.value.internet == "true") |
        select(.key | test(":")) |
        .key
    ' "$POLICY_FILE" 2>/dev/null)

    # Log + drop (final rule)
    cat >> "$NFT_FILE" << EOF

        /* log + drop */
        iifname "wg0" limit rate ${LOG_RATE} log prefix "${LOG_PREFIX}: " drop
    }

    chain postrouting {
        type nat hook postrouting priority srcnat; policy accept;
        oifname "${DEF_IF}" masquerade
    }
}
EOF
}

# ============================================================
# 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 current nftables state
    backup_nftables

    # 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 ===
    sysctl -w net.ipv4.ip_forward=1 >/dev/null 2>&1 || true
    sysctl -w net.ipv6.conf.all.forwarding=1 >/dev/null 2>&1 || true

    # === BACKUP EXISTING TABLE ===
    if nft list table "$NFT_TABLE_FULL" &>/dev/null; then
        nft list table "$NFT_TABLE_FULL" > "$NFT_BACKUP" 2>/dev/null || true
        log_info "Backed up existing table to $NFT_BACKUP"
    fi

    # === DELETE OLD TABLE ===
    nft delete table "$NFT_TABLE_FULL" 2>/dev/null || true

    # === GENERATE AND LOAD NEW RULESET ===
    generate_ruleset "$WG_SUBNET" "$WG_SUBNET_V6" "$LAN_SUBNETS" "$DEF_IF"
    log_info "Generated ruleset: $NFT_FILE"

    nft -f "$NFT_FILE"
    log_info "Loaded nftables ruleset from $NFT_FILE"

    # === POPULATE SETS ===
    log_info "Populating whitelist sets..."

    local v4_count=0
    local v6_count=0

    while IFS= read -r line; do
        [[ -z "$line" ]] && continue
        local client_ip target
        client_ip=$(echo "$line" | awk '{print $1}')
        target=$(echo "$line" | awk '{print $2}')
        [[ -z "$client_ip" || -z "$target" ]] && continue

        if [[ "$target" == *":"* ]]; then
            nft add element "$NFT_TABLE_FULL" "$NFT_SET_V6" { "$client_ip" . "$target" } 2>/dev/null || \
                log_warn "Failed to add ${client_ip} . ${target} to set $NFT_SET_V6"
            (( v6_count++ )) || true
        else
            nft add element "$NFT_TABLE_FULL" "$NFT_SET_V4" { "$client_ip" . "$target" } 2>/dev/null || \
                log_warn "Failed to add ${client_ip} . ${target} to set $NFT_SET_V4"
            (( v4_count++ )) || true
        fi
    done < <(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)

    log_info "Set $NFT_SET_V4: $v4_count entries, $NFT_SET_V6: $v6_count entries"

    # === CLEANUP BACKUP (success path) ===
    rm -f "$NFT_BACKUP" 2>/dev/null || true

    # Disable ERR trap (success path)
    trap - ERR

    # === VERIFY ===
    local rule_count
    rule_count=$(nft list chain "$NFT_TABLE_FULL" forward 2>/dev/null | grep -c '^\s*' || echo 0)
    log_info "Policy applied. Table: $NFT_TABLE, Rules: $rule_count"

    echo "[OK] nftables policy applied. Table: $NFT_TABLE"
}

main "$@"