wireguard-vpn/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 NFT_TABLE="wg_policy"
readonly NFT_TABLE_FULL="inet wg_policy"
readonly NFT_SET_V4="wg_allowed_v4"
readonly NFT_SET_V6="wg_allowed_v6"
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/minute"
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() {
    local ip="$1"
    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
        if [[ "$octet" =~ ^0[0-9] ]]; then
            return 1
        fi
        if (( octet < 0 || octet > 255 )); then
            return 1
        fi
    done
    return 0
}

validate_ipv4_cidr() {
    local cidr="$1"
    local ip prefix

    if [[ "$cidr" == *"/"* ]]; then
        ip="${cidr%%/*}"
        prefix="${cidr##*/}"
    else
        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() {
    local ip="$1"
    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() {
    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
}

validate_cidr() {
    local cidr="$1"
    if [[ "$cidr" == *":"* ]]; then
        validate_ipv6_cidr "$cidr"
    else
        validate_ipv4_cidr "$cidr"
    fi
}

# ============================================================
# NFTABLES HELPERS
# ============================================================

has_nft() {
    command -v nft &>/dev/null
}

# ============================================================
# 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 ))
    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() {
    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

    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_nftables() {
    mkdir -p "$BACKUP_DIR"
    local timestamp
    timestamp="$(date '+%Y%m%d_%H%M%S')"

    if nft list ruleset > "${BACKUP_DIR}/nftables_${timestamp}.rules" 2>/dev/null; then
        log_info "nftables backup: ${BACKUP_DIR}/nftables_${timestamp}.rules"
    fi

    local count
    count=$(find "$BACKUP_DIR" -name 'nftables_*.rules' -type f | wc -l)
    if (( count > 20 )); then
        find "$BACKUP_DIR" -name 'nftables_*.rules' -type f -printf '%T@ %p\n' \
            | sort -n \
            | head -n $(( count - 20 )) \
            | awk '{print $2}' \
            | xargs rm -f
    fi
}

# ============================================================
# 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 nftables table exists
    if nft list table "$NFT_TABLE_FULL" &>/dev/null; then
        local rule_count
        rule_count=$(nft list chain "$NFT_TABLE_FULL" forward 2>/dev/null | grep -c '^\s*' || echo 0)
        report+="[OK] Table $NFT_TABLE active ($rule_count rules)\n"
    else
        report+="[WARN] Table $NFT_TABLE not found\n"
        status=1
    fi

    # 4. Check forward chain exists in table
    if nft list chain "$NFT_TABLE_FULL" forward &>/dev/null; then
        report+="[OK] Forward chain exists in $NFT_TABLE\n"
    else
        report+="[WARN] Forward chain not found in $NFT_TABLE\n"
        status=1
    fi

    # 5. Check nftables sets
    for set_name in "$NFT_SET_V4" "$NFT_SET_V6"; do
        if nft list set "$NFT_TABLE_FULL" "$set_name" &>/dev/null; then
            local entry_count
            entry_count=$(nft list set "$NFT_TABLE_FULL" "$set_name" 2>/dev/null | grep -c '^\s*' || echo 0)
            report+="[OK] Set $set_name active ($entry_count entries)\n"
        else
            report+="[INFO] Set $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
}