#!/bin/bash
set -euo pipefail

# =============================================================================
# Arvandor Port Forwarding Script
# =============================================================================
# Configures NAT (DNAT/SNAT) and FORWARD rules for Proxmox host.
# Uses a custom chain (ARVANDOR-FORWARD) to avoid conflicts with PVE firewall.
#
# Usage:
#   ./port-forward.sh           # Apply rules
#   ./port-forward.sh --dry-run # Show what would be done
#   ./port-forward.sh --restore # Restore backup
#   ./port-forward.sh --status  # Show current rules
# =============================================================================

# -----------------------------------------------------------------------------
# Configuration - UPDATE THESE FOR YOUR ENVIRONMENT
# -----------------------------------------------------------------------------
NETWORK_INTERFACE="vmbr0"
INTERNAL_NETWORK="192.168.100.0/24"
PUBLIC_IP="203.0.113.10"  # Your public IP
CUSTOM_CHAIN="ARVANDOR-FORWARD"
BACKUP_FILE="/root/network/iptables.backup"

# Nebula Lighthouse
NEBULA_IP="192.168.100.10"
NEBULA_PORT="4242"

# Caddy (Reverse Proxy)
CADDY_IP="192.168.100.12"
CADDY_HTTP_PORT="80"
CADDY_HTTPS_PORT="443"

# Gitea (Optional)
GITEA_IP="192.168.100.23"
GITEA_SSH_PORT="2222"

# Security - restrict SSH to specific IP
ALLOWED_SSH_IP="203.0.113.20"  # Your home IP

# -----------------------------------------------------------------------------
# Functions
# -----------------------------------------------------------------------------
log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*"
}

error() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] ERROR: $*" >&2
}

die() {
    error "$*"
    exit 1
}

check_root() {
    [[ $EUID -eq 0 ]] || die "This script must be run as root"
}

check_interface() {
    local iface=$1
    ip link show "$iface" &>/dev/null || die "Interface $iface does not exist"
}

backup_rules() {
    log "Backing up current iptables rules to $BACKUP_FILE"
    mkdir -p "$(dirname "$BACKUP_FILE")"
    iptables-save > "$BACKUP_FILE"
}

restore_rules() {
    [[ -f "$BACKUP_FILE" ]] || die "Backup file $BACKUP_FILE not found"
    log "Restoring iptables rules from $BACKUP_FILE"
    iptables-restore < "$BACKUP_FILE"
    log "Rules restored successfully"
}

setup_custom_chain() {
    # Create custom chain if it doesn't exist
    if ! iptables -L "$CUSTOM_CHAIN" -n &>/dev/null; then
        log "Creating custom chain: $CUSTOM_CHAIN"
        iptables -N "$CUSTOM_CHAIN"
    fi

    # Ensure chain is jumped to from FORWARD (only once)
    if ! iptables -C FORWARD -j "$CUSTOM_CHAIN" &>/dev/null; then
        log "Inserting jump to $CUSTOM_CHAIN in FORWARD chain"
        iptables -I FORWARD 1 -j "$CUSTOM_CHAIN"
    fi

    # Flush the custom chain
    log "Flushing custom chain: $CUSTOM_CHAIN"
    iptables -F "$CUSTOM_CHAIN"
}

apply_rules() {
    local dry_run=${1:-false}

    if [[ "$dry_run" == "true" ]]; then
        log "=== DRY RUN MODE - No changes will be made ==="
        echo ""
        echo "Would apply the following rules:"
        echo ""
        echo "NAT PREROUTING (DNAT):"
        echo "  - UDP $NEBULA_PORT → $NEBULA_IP:$NEBULA_PORT (Nebula)"
        echo "  - TCP $CADDY_HTTP_PORT → $CADDY_IP:$CADDY_HTTP_PORT (HTTP)"
        echo "  - TCP $CADDY_HTTPS_PORT → $CADDY_IP:$CADDY_HTTPS_PORT (HTTPS)"
        echo "  - TCP $GITEA_SSH_PORT → $GITEA_IP:$GITEA_SSH_PORT (Gitea SSH)"
        echo ""
        echo "FORWARD chain ($CUSTOM_CHAIN):"
        echo "  - Allow traffic to all above destinations"
        echo ""
        echo "INPUT:"
        echo "  - Allow Nebula (nebula1 interface)"
        echo "  - Allow SSH from $ALLOWED_SSH_IP"
        echo "  - Drop SSH from all others"
        echo "  - Block Proxmox UI from $NETWORK_INTERFACE"
        return
    fi

    # --- NAT Rules ---
    log "Flushing NAT rules..."
    iptables -t nat -F PREROUTING
    iptables -t nat -F POSTROUTING

    log "Setting up NAT masquerading..."
    iptables -t nat -A POSTROUTING -s "$INTERNAL_NETWORK" -o "$NETWORK_INTERFACE" -j MASQUERADE

    log "Setting up hairpin NAT for Nebula..."
    iptables -t nat -A PREROUTING -s "$INTERNAL_NETWORK" -d "$PUBLIC_IP" -p udp --dport "$NEBULA_PORT" -j DNAT --to-destination "$NEBULA_IP:$NEBULA_PORT"
    iptables -t nat -A POSTROUTING -s "$INTERNAL_NETWORK" -d "$NEBULA_IP" -p udp --dport "$NEBULA_PORT" -j SNAT --to-source "$PUBLIC_IP"

    log "Setting up hairpin NAT for Gitea SSH..."
    iptables -t nat -A PREROUTING -s "$INTERNAL_NETWORK" -d "$PUBLIC_IP" -p tcp --dport "$GITEA_SSH_PORT" -j DNAT --to-destination "$GITEA_IP:$GITEA_SSH_PORT"
    iptables -t nat -A POSTROUTING -s "$INTERNAL_NETWORK" -d "$GITEA_IP" -p tcp --dport "$GITEA_SSH_PORT" -j SNAT --to-source "$PUBLIC_IP"

    log "Setting up DNAT rules..."
    # Nebula
    iptables -t nat -A PREROUTING -i "$NETWORK_INTERFACE" -p udp --dport "$NEBULA_PORT" -j DNAT --to-destination "$NEBULA_IP:$NEBULA_PORT"
    # Caddy
    iptables -t nat -A PREROUTING -i "$NETWORK_INTERFACE" -p tcp --dport "$CADDY_HTTP_PORT" -j DNAT --to-destination "$CADDY_IP:$CADDY_HTTP_PORT"
    iptables -t nat -A PREROUTING -i "$NETWORK_INTERFACE" -p tcp --dport "$CADDY_HTTPS_PORT" -j DNAT --to-destination "$CADDY_IP:$CADDY_HTTPS_PORT"
    # Gitea SSH
    iptables -t nat -A PREROUTING -i "$NETWORK_INTERFACE" -p tcp --dport "$GITEA_SSH_PORT" -j DNAT --to-destination "$GITEA_IP:$GITEA_SSH_PORT"

    # --- FORWARD Rules (custom chain) ---
    setup_custom_chain

    log "Adding FORWARD rules to $CUSTOM_CHAIN..."
    iptables -A "$CUSTOM_CHAIN" -d "$CADDY_IP" -p tcp --dport "$CADDY_HTTP_PORT" -j ACCEPT
    iptables -A "$CUSTOM_CHAIN" -d "$CADDY_IP" -p tcp --dport "$CADDY_HTTPS_PORT" -j ACCEPT
    iptables -A "$CUSTOM_CHAIN" -d "$NEBULA_IP" -p udp --dport "$NEBULA_PORT" -j ACCEPT
    iptables -A "$CUSTOM_CHAIN" -d "$GITEA_IP" -p tcp --dport "$GITEA_SSH_PORT" -j ACCEPT

    # --- INPUT Rules ---
    log "Flushing INPUT rules..."
    iptables -F INPUT

    log "Setting up INPUT rules..."
    iptables -A INPUT -i nebula1 -j ACCEPT
    iptables -A INPUT -p tcp --dport 22 -s "$ALLOWED_SSH_IP" -j ACCEPT
    iptables -A INPUT -p tcp --dport 22 -j DROP
    iptables -I INPUT -i "$NETWORK_INTERFACE" -p tcp --dport 8006 -j DROP
    iptables -I INPUT -i vmbr1 -p tcp --dport 8006 -j ACCEPT
}

save_rules() {
    log "Saving iptables rules persistently..."
    if command -v netfilter-persistent &>/dev/null; then
        netfilter-persistent save
        log "Rules saved via netfilter-persistent"
    else
        die "netfilter-persistent not found. Install with: apt install iptables-persistent"
    fi
}

show_status() {
    echo ""
    echo "=== Port Forwarding Status ==="
    echo ""
    echo "NAT PREROUTING rules:"
    iptables -t nat -L PREROUTING -n --line-numbers 2>/dev/null | head -20
    echo ""
    echo "FORWARD chain ($CUSTOM_CHAIN):"
    iptables -L "$CUSTOM_CHAIN" -n --line-numbers 2>/dev/null || echo "Chain not found"
    echo ""
    echo "=== Services ==="
    echo "  HTTP/HTTPS: 80,443 → Caddy ($CADDY_IP)"
    echo "  Nebula:     $NEBULA_PORT → Lighthouse ($NEBULA_IP)"
    echo "  Gitea SSH:  $GITEA_SSH_PORT → $GITEA_IP"
}

# -----------------------------------------------------------------------------
# Main
# -----------------------------------------------------------------------------
main() {
    local action="${1:-apply}"

    case "$action" in
        --dry-run|-n)
            check_root
            check_interface "$NETWORK_INTERFACE"
            apply_rules true
            ;;
        --restore|-r)
            check_root
            restore_rules
            ;;
        --status|-s)
            show_status
            ;;
        apply|"")
            check_root
            check_interface "$NETWORK_INTERFACE"
            backup_rules
            apply_rules false
            save_rules
            log "Setup complete!"
            show_status
            ;;
        *)
            echo "Usage: $0 [--dry-run|--restore|--status]"
            exit 1
            ;;
    esac
}

main "$@"