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