#!/bin/bash
# Manage port forwarding from VPS to home server via WireGuard

set -e

INTERFACE=$(ip route | grep default | awk '{print $5}' | head -n1)
HOME_SERVER="10.0.0.2"

show_help() {
    echo "Port Forwarding Manager"
    echo "========================"
    echo "Forward ports from VPS to home server via WireGuard tunnel."
    echo ""
    echo "Usage: $0 <command> [options]"
    echo ""
    echo "Commands:"
    echo "  list                      Show all currently forwarded ports"
    echo "  add <port> [protocol]     Forward a port to home server"
    echo "  remove <port> [protocol]  Stop forwarding a port"
    echo "  help                      Show this help message"
    echo ""
    echo "Protocol options:"
    echo "  tcp   - TCP only (default)"
    echo "  udp   - UDP only"
    echo "  both  - TCP and UDP"
    echo ""
    echo "Examples:"
    echo "  $0 list                   # Show all forwarded ports"
    echo "  $0 add 25565              # Forward Minecraft server (TCP)"
    echo "  $0 add 7777 udp           # Forward UDP port 7777"
    echo "  $0 add 27015 both         # Forward game server (TCP+UDP)"
    echo "  $0 remove 25565           # Stop forwarding port 25565"
    echo "  $0 remove 7777 udp        # Stop forwarding UDP port 7777"
    echo ""
    echo "Notes:"
    echo "  - Ports 80 and 443 are reserved for Traefik (HTTP/HTTPS)"
    echo "  - Changes are saved automatically and persist after reboot"
    echo "  - Run as root (sudo)"
}

show_add_usage() {
    echo "Error: Missing port number"
    echo ""
    echo "Usage: $0 add <port> [protocol]"
    echo ""
    echo "Arguments:"
    echo "  port      Port number to forward (required)"
    echo "  protocol  tcp (default), udp, or both"
    echo ""
    echo "Examples:"
    echo "  $0 add 25565        # Forward TCP port 25565"
    echo "  $0 add 7777 udp     # Forward UDP port 7777"
    echo "  $0 add 27015 both   # Forward both TCP and UDP"
}

show_remove_usage() {
    echo "Error: Missing port number"
    echo ""
    echo "Usage: $0 remove <port> [protocol]"
    echo ""
    echo "Arguments:"
    echo "  port      Port number to stop forwarding (required)"
    echo "  protocol  tcp (default), udp, or both"
    echo ""
    echo "Examples:"
    echo "  $0 remove 25565        # Stop forwarding TCP port 25565"
    echo "  $0 remove 7777 udp     # Stop forwarding UDP port 7777"
    echo "  $0 remove 27015 both   # Stop forwarding both TCP and UDP"
}

show_invalid_protocol() {
    echo "Error: Invalid protocol '$1'"
    echo ""
    echo "Valid protocols:"
    echo "  tcp   - TCP only (default)"
    echo "  udp   - UDP only"
    echo "  both  - TCP and UDP"
}

list_ports() {
    echo "=== Forwarded Ports ==="
    echo ""
    local found=0
    iptables -t nat -L PREROUTING -n 2>/dev/null | grep -E "dpt:" | grep "10.0.0.2" | \
        awk '{for(i=1;i<=NF;i++) if($i ~ /dpt:/) print $2, $i}' | \
        sed 's/dpt://' | \
        while read proto port; do
            echo "  $port ($proto)"
            found=1
        done

    if [ $found -eq 0 ]; then
        # Check again since subshell doesn't preserve variable
        if ! iptables -t nat -L PREROUTING -n 2>/dev/null | grep -q "10.0.0.2.*dpt:"; then
            echo "  No ports currently forwarded (besides 80/443 for Traefik)"
        fi
    fi
    echo ""
}

add_port() {
    local port=$1
    local proto=$2

    echo "Adding $proto port $port -> $HOME_SERVER:$port"
    iptables -t nat -A PREROUTING -i $INTERFACE -p $proto --dport $port -j DNAT --to-destination $HOME_SERVER:$port
    iptables -A FORWARD -i $INTERFACE -o wg0 -p $proto --dport $port -j ACCEPT
}

remove_port() {
    local port=$1
    local proto=$2

    echo "Removing $proto port $port"
    iptables -t nat -D PREROUTING -i $INTERFACE -p $proto --dport $port -j DNAT --to-destination $HOME_SERVER:$port 2>/dev/null || echo "  NAT rule not found"
    iptables -D FORWARD -i $INTERFACE -o wg0 -p $proto --dport $port -j ACCEPT 2>/dev/null || echo "  FORWARD rule not found"
}

save_rules() {
    echo ""
    echo "Saving iptables rules..."
    netfilter-persistent save >/dev/null 2>&1
    echo "Done! Changes will persist after reboot."
}

# Main
case ${1:-} in
    help|--help|-h)
        show_help
        ;;
    list)
        list_ports
        ;;
    add)
        if [ -z "${2:-}" ]; then
            show_add_usage
            exit 1
        fi
        PORT=$2
        PROTOCOL=${3:-tcp}

        case $PROTOCOL in
            tcp)
                add_port $PORT tcp
                ;;
            udp)
                add_port $PORT udp
                ;;
            both)
                add_port $PORT tcp
                add_port $PORT udp
                ;;
            *)
                show_invalid_protocol $PROTOCOL
                exit 1
                ;;
        esac
        save_rules
        ;;
    remove)
        if [ -z "${2:-}" ]; then
            show_remove_usage
            exit 1
        fi
        PORT=$2
        PROTOCOL=${3:-tcp}

        case $PROTOCOL in
            tcp)
                remove_port $PORT tcp
                ;;
            udp)
                remove_port $PORT udp
                ;;
            both)
                remove_port $PORT tcp
                remove_port $PORT udp
                ;;
            *)
                show_invalid_protocol $PROTOCOL
                exit 1
                ;;
        esac
        save_rules
        ;;
    "")
        echo "Error: No command specified"
        echo ""
        echo "Usage: $0 <command> [options]"
        echo ""
        echo "Commands:"
        echo "  list                      Show all forwarded ports"
        echo "  add <port> [protocol]     Forward a port"
        echo "  remove <port> [protocol]  Stop forwarding a port"
        echo "  help                      Show detailed help"
        echo ""
        echo "Run '$0 help' for more information."
        exit 1
        ;;
    *)
        echo "Error: Unknown command '$1'"
        echo ""
        echo "Valid commands: list, add, remove, help"
        echo ""
        echo "Run '$0 help' for more information."
        exit 1
        ;;
esac