#!/bin/bash

# Colors and formatting
RED='\033[0;31m'
GREEN='\033[0;32m'
BLUE='\033[0;34m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
NC='\033[0m' # No Color
BOLD='\033[1m'

# Default essential ports
DEFAULT_TCP_PORT=22
DEFAULT_UDP_PORT=41641

# Check if running as root
if [ "$EUID" -ne 0 ]; then
    echo -e "${RED}${BOLD}[ERROR]${NC} Please run as root"
    exit 1
fi

# Check for iptables-persistent installation
check_iptables_persistent() {
    if ! dpkg -l | grep -q "iptables-persistent"; then
        echo -e "${YELLOW}[INFO]${NC} iptables-persistent is not installed"
        read -p $'\033[0;36mWould you like to install iptables-persistent? [Y/n]: \033[0m' install_choice
        install_choice=${install_choice:-Y}
        
        if [[ $install_choice =~ ^[Yy]$ ]]; then
            # Pre-seed the debconf database to avoid interactive prompt
            echo iptables-persistent iptables-persistent/autosave_v4 boolean true | debconf-set-selections
            echo iptables-persistent iptables-persistent/autosave_v6 boolean true | debconf-set-selections
            
            apt-get update
            if apt-get install -y iptables-persistent; then
                echo -e "${GREEN}${BOLD}[SUCCESS]${NC} iptables-persistent installed successfully"
            else
                echo -e "${RED}${BOLD}[ERROR]${NC} Failed to install iptables-persistent"
                return 1
            fi
        else
            echo -e "${YELLOW}[WARNING]${NC} Continuing without iptables-persistent"
            return 1
        fi
    fi
    return 0
}

# Enhanced save function with iptables-persistent support
save_iptables() {
    echo -e "${YELLOW}[INFO]${NC} Saving iptables rules..."
    
    # Create backup directory if it doesn't exist
    mkdir -p /etc/iptables/backup
    
    # Create timestamped backup
    timestamp=$(date +%Y%m%d_%H%M%S)
    backup_file="/etc/iptables/backup/rules.v4.${timestamp}"
    
    # Backup current rules
    if ! iptables-save > "$backup_file"; then
        echo -e "${RED}${BOLD}[ERROR]${NC} Failed to create backup"
        return 1
    fi
    
    # Try saving with netfilter-persistent first
    if command -v netfilter-persistent &>/dev/null; then
        if netfilter-persistent save; then
            echo -e "${GREEN}${BOLD}[SUCCESS]${NC} Rules saved using netfilter-persistent"
            echo -e "${GREEN}[INFO]${NC} Backup saved to: ${backup_file}"
            return 0
        fi
    fi
    
    # Fallback to manual save
    if [ -d "/etc/iptables" ]; then
        if iptables-save > /etc/iptables/rules.v4; then
            echo -e "${GREEN}${BOLD}[SUCCESS]${NC} Rules saved to /etc/iptables/rules.v4"
            echo -e "${GREEN}[INFO]${NC} Backup saved to: ${backup_file}"
            return 0
        fi
    fi
    
    echo -e "${RED}${BOLD}[ERROR]${NC} Failed to save iptables rules"
    return 1
}

# Enhanced restore function with iptables-persistent support
restore_iptables() {
    echo -e "${YELLOW}[INFO]${NC} Attempting to restore previous rules..."
    
    # Try restoring from backup first
    local latest_backup=$(ls -t /etc/iptables/backup/rules.v4.* 2>/dev/null | head -n1)
    
    if [ -n "$latest_backup" ]; then
        echo -e "${YELLOW}[INFO]${NC} Restoring from latest backup: ${latest_backup}"
        if iptables-restore < "$latest_backup"; then
            echo -e "${GREEN}${BOLD}[SUCCESS]${NC} Rules restored from backup"
            # Save the restored rules
            save_iptables
            return 0
        fi
    fi
    
    # Fallback to default rules file
    if [ -f /etc/iptables/rules.v4 ]; then
        echo -e "${YELLOW}[INFO]${NC} Attempting to restore from /etc/iptables/rules.v4"
        if iptables-restore < /etc/iptables/rules.v4; then
            echo -e "${GREEN}${BOLD}[SUCCESS]${NC} Rules restored from default file"
            return 0
        fi
    fi
    
    echo -e "${RED}${BOLD}[ERROR]${NC} Failed to restore iptables rules"
    return 1
}

# Add this near the beginning of the script execution
check_iptables_persistent

# Function to print section header
print_header() {
    echo -e "\n${BLUE}${BOLD}$1${NC}"
    echo -e "${BLUE}${BOLD}$(printf '=%.0s' {1..50})${NC}\n"
}

# Function to display current NAT configuration
display_current_config() {
    print_header "Current NAT Configuration"
    
    echo -e "${CYAN}PREROUTING Rules (DNAT):${NC}"
    if ! iptables -t nat -L PREROUTING -n -v | grep -q "DNAT"; then
        echo -e "${YELLOW}No DNAT rules configured${NC}"
    else
        iptables -t nat -L PREROUTING -n -v | grep "DNAT"
    fi
    
    echo -e "\n${CYAN}POSTROUTING Rules (SNAT/MASQUERADE):${NC}"
    if ! iptables -t nat -L POSTROUTING -n -v | grep -q "SNAT\|MASQUERADE"; then
        echo -e "${YELLOW}No SNAT/MASQUERADE rules configured${NC}"
    else
        iptables -t nat -L POSTROUTING -n -v | grep "SNAT\|MASQUERADE"
    fi
    
    echo -e "\n${CYAN}Forward Rules:${NC}"
    if ! iptables -L FORWARD -n -v | grep -q "ACCEPT\|DROP"; then
        echo -e "${YELLOW}No specific forwarding rules configured${NC}"
    else
        iptables -L FORWARD -n -v | grep "ACCEPT\|DROP"
    fi
    
    echo # Empty line for spacing
}

# Function to get IP and interface pairs (excluding localhost)
get_interfaces() {
    ip -4 addr | awk '/inet / && !/127.0.0.1/ {print $2 "\t" $NF}' | sed 's/\/[0-9]*//g'
}

# Get interface name from selected IP
get_interface_name() {
    local ip=$1
    ip -4 addr | grep "$ip" | awk '{print $NF}'
}

# Function to get Tailscale IPs
get_tailscale_devices() {
    tailscale status | grep -v '^#' | awk 'NF>1 {print $2 " - " $1}'
}

# Function to get interface details for Tailscale device
get_tailscale_device_name() {
    local ip=$1
    tailscale status | grep "$ip" | awk '{print $2}'
}

# Function to get subnet from IP
get_subnet() {
    local ip=$1
    ip -4 addr | grep "$ip" | awk '{print $2}'
}

# Display current configuration at startup
display_current_config

# Step 1: Get source IP
print_header "Step 1: Select source IP address"
echo -e "${CYAN}Available interfaces:${NC}"
echo -e "${YELLOW}0)${NC} Enter custom IP"

# Display interface options
declare -a interfaces
while IFS=$'\t' read -r ip interface; do
    interfaces+=("$ip - $interface")
done < <(get_interfaces)

for i in "${!interfaces[@]}"; do
    echo -e "${YELLOW}$((i + 1)))${NC} ${interfaces[$i]}"
done
echo -e "${YELLOW}C)${NC} Cancel"
echo -e "${CYAN}Press Enter to select option 1${NC}"

read -p $'\033[0;36mSelect an option: \033[0m' choice

# Default to option 1 if Enter is pressed
choice=${choice:-1}

case $choice in
0)
    read -p $'\033[0;36mEnter custom IP: \033[0m' source_ip
    source_subnet=$(get_subnet "$source_ip")
    ;;
[1-9]*)
    if [ "$choice" -le "${#interfaces[@]}" ]; then
        source_ip=$(echo "${interfaces[$((choice - 1))]}" | cut -d' ' -f1)
        source_subnet=$(get_subnet "$source_ip")
    else
        echo -e "${RED}${BOLD}[ERROR]${NC} Invalid option"
        exit 1
    fi
    ;;
[Cc])
    echo -e "${YELLOW}Operation cancelled${NC}"
    exit 0
    ;;
*)
    echo -e "${RED}${BOLD}[ERROR]${NC} Invalid option"
    exit 1
    ;;
esac

# Step 2: Get target IP
print_header "Step 2: Select target Tailscale IP address"
echo -e "${CYAN}Available Tailscale devices:${NC}"
echo -e "${YELLOW}0)${NC} Enter custom IP"

# Display Tailscale devices
declare -a tailscale_devices
mapfile -t tailscale_devices < <(get_tailscale_devices)

for i in "${!tailscale_devices[@]}"; do
    echo -e "${YELLOW}$((i + 1)))${NC} ${tailscale_devices[$i]}"
done
echo -e "${YELLOW}C)${NC} Cancel"
echo -e "${CYAN}Press Enter to select option 0 (custom IP)${NC}"

read -p $'\033[0;36mSelect an option: \033[0m' choice

# Default to option 0 if Enter is pressed
choice=${choice:-0}

case $choice in
0)
    read -p $'\033[0;36mEnter custom IP: \033[0m' target_ip
    target_device="custom"
    ;;
[1-9]*)
    if [ "$choice" -le "${#tailscale_devices[@]}" ]; then
        selected_device="${tailscale_devices[$((choice-1))]}"
        target_ip=$(echo "$selected_device" | cut -d' ' -f3)
        target_device=$(echo "$selected_device" | cut -d' ' -f1)
    else
        echo -e "${RED}${BOLD}[ERROR]${NC} Invalid option"
        exit 1
    fi
    ;;
[Cc])
    echo -e "${YELLOW}Operation cancelled${NC}"
    exit 0
    ;;
*)
    echo -e "${RED}${BOLD}[ERROR]${NC} Invalid option"
    exit 1
    ;;
esac

# Step 3: Port Configuration
print_header "Step 3: Configure Port Settings"
echo -e "${CYAN}Port Configuration Options:${NC}"
echo -e "${YELLOW}1)${NC} Use default ports only (SSH TCP 22, Tailscale UDP 41641)"
echo -e "${YELLOW}2)${NC} Configure custom ports"
echo -e "${CYAN}Press Enter to use default ports only${NC}"

read -p $'\033[0;36mSelect an option: \033[0m' port_choice
port_choice=${port_choice:-1}

declare -a PRESERVED_TCP_PORTS
declare -a PRESERVED_UDP_PORTS

case $port_choice in
1)
    PRESERVED_TCP_PORTS=($DEFAULT_TCP_PORT)
    PRESERVED_UDP_PORTS=($DEFAULT_UDP_PORT)
    echo -e "${GREEN}Using default TCP port: ${DEFAULT_TCP_PORT}${NC}"
    echo -e "${GREEN}Using default UDP port: ${DEFAULT_UDP_PORT}${NC}"
    ;;
2)
    echo -e "\n${CYAN}Enter TCP ports to preserve (space-separated, e.g., '22 80 443')${NC}"
    echo -e "${CYAN}Press Enter to use default TCP port ${DEFAULT_TCP_PORT}${NC}"
    read -p $'\033[0;36mTCP Ports: \033[0m' custom_tcp_ports
    
    if [ -z "$custom_tcp_ports" ]; then
        PRESERVED_TCP_PORTS=($DEFAULT_TCP_PORT)
        echo -e "${GREEN}Using default TCP port: ${DEFAULT_TCP_PORT}${NC}"
    else
        PRESERVED_TCP_PORTS=($custom_tcp_ports)
        echo -e "${GREEN}Using custom TCP ports: ${custom_tcp_ports}${NC}"
    fi

    echo -e "\n${CYAN}Enter UDP ports to preserve (space-separated, e.g., '41641 53')${NC}"
    echo -e "${CYAN}Press Enter to use default UDP port ${DEFAULT_UDP_PORT}${NC}"
    read -p $'\033[0;36mUDP Ports: \033[0m' custom_udp_ports
    
    if [ -z "$custom_udp_ports" ]; then
        PRESERVED_UDP_PORTS=($DEFAULT_UDP_PORT)
        echo -e "${GREEN}Using default UDP port: ${DEFAULT_UDP_PORT}${NC}"
    else
        PRESERVED_UDP_PORTS=($custom_udp_ports)
        echo -e "${GREEN}Using custom UDP ports: ${custom_udp_ports}${NC}"
    fi
    ;;
*)
    echo -e "${RED}${BOLD}[ERROR]${NC} Invalid option"
    exit 1
    ;;
esac

# Step 4: Confirmation
print_header "Step 4: Review your selections"
source_interface=$(get_interface_name "$source_ip")
echo -e "${CYAN}Source Details:${NC}"
echo -e "  IP Address: ${BOLD}$source_ip${NC}"
echo -e "  Interface:  ${BOLD}$source_interface${NC}"
echo -e "  Subnet:     ${BOLD}$source_subnet${NC}"
echo -e "\n${CYAN}Target Details:${NC}"
echo -e "  IP Address: ${BOLD}$target_ip${NC}"
echo -e "  Device:     ${BOLD}$target_device${NC}"
echo -e "\n${CYAN}Port Configuration:${NC}"
echo -e "  Preserved TCP Ports: ${BOLD}${PRESERVED_TCP_PORTS[*]}${NC}"
echo -e "  Preserved UDP Ports: ${BOLD}${PRESERVED_UDP_PORTS[*]}${NC}"
echo -e "\n${YELLOW}This will forward traffic from $source_interface to tailscale0${NC}"
read -p $'\033[1;33mApply these changes? [Y/N, default is Y]: \033[0m' confirm

# Default to 'Y' if no input is provided
confirm=${confirm:-Y}

if [[ $confirm =~ ^[Yy]$ ]]; then
    # Step 5: Apply iptables configuration
    print_header "Step 5: Applying configuration..."

    # Create backup before any changes
    echo -e "${YELLOW}[1/7]${NC} Creating backup..."
    mkdir -p /etc/iptables/backup
    timestamp=$(date +%Y%m%d_%H%M%S)
    backup_file="/etc/iptables/backup/rules.v4.${timestamp}"

    if ! iptables-save > "$backup_file"; then
        echo -e "${RED}${BOLD}[ERROR]${NC} Failed to create backup. Aborting configuration"
        exit 1
    fi

    if [ ! -s "$backup_file" ]; then
        echo -e "${RED}${BOLD}[ERROR]${NC} Backup file is empty. Aborting configuration"
        exit 1
    fi

    # Check if IP forwarding is enabled
    if [ "$(sysctl -n net.ipv4.ip_forward)" -ne 1 ]; then
        echo -e "${YELLOW}[2/7]${NC} Enabling IP forwarding..."
        if ! echo "net.ipv4.ip_forward=1" >> /etc/sysctl.conf || ! sysctl -w net.ipv4.ip_forward=1; then
            echo -e "${RED}${BOLD}[ERROR]${NC} Failed to enable IP forwarding"
            exit 1
        fi
        # Also add to sysctl.d for persistent configuration
        if ! echo "net.ipv4.ip_forward=1" > /etc/sysctl.d/99-tailscale.conf; then
            echo -e "${RED}${BOLD}[ERROR]${NC} Failed to create persistent IP forwarding configuration"
            exit 1
        fi
        sysctl -p /etc/sysctl.d/99-tailscale.conf
    else
        echo -e "${YELLOW}[2/7]${NC} IP forwarding is already enabled"
    fi

    # Get interface name
    source_interface=$(get_interface_name "$source_ip")
    if [ -z "$source_interface" ]; then
        echo -e "${RED}${BOLD}[ERROR]${NC} Could not determine interface name for IP $source_ip"
        exit 1
    fi

    # Apply rules
    echo -e "${YELLOW}[3/7]${NC} Applying new rules..."
    
    # Clear existing NAT rules
    echo -e "${CYAN}Clearing existing NAT and forwarding rules...${NC}"
    iptables -t nat -F
    iptables -F FORWARD

    # Allow preserved TCP ports to reach the host
    for port in "${PRESERVED_TCP_PORTS[@]}"; do
        echo -e "${CYAN}Preserving TCP host access on port ${port}${NC}"
        iptables -t nat -A PREROUTING -d "${source_ip}" -p tcp --dport ${port} -j ACCEPT
        iptables -A FORWARD -p tcp --dport ${port} -j ACCEPT
    done

    # Allow preserved UDP ports to reach the host
    for port in "${PRESERVED_UDP_PORTS[@]}"; do
        echo -e "${CYAN}Preserving UDP host access on port ${port}${NC}"
        iptables -t nat -A PREROUTING -d "${source_ip}" -p udp --dport ${port} -j ACCEPT
        iptables -A FORWARD -p udp --dport ${port} -j ACCEPT
        # Additional rule to ensure UDP responses can return
        iptables -A FORWARD -p udp --sport ${port} -j ACCEPT
    done

    # Set up NAT rules for all other ports
    echo -e "${CYAN}Setting up NAT rules...${NC}"
    if ! iptables -t nat -A PREROUTING -d "${source_ip}" -j DNAT --to-destination "${target_ip}"; then
        echo -e "${RED}${BOLD}[ERROR]${NC} Failed to set up DNAT rules"
        echo -e "${YELLOW}[INFO]${NC} Attempting to restore previous configuration..."
        if ! iptables-restore < "$backup_file"; then
            echo -e "${RED}${BOLD}[ERROR]${NC} Failed to restore from backup: $backup_file"
            echo -e "${RED}${BOLD}[WARNING]${NC} Your firewall rules may be in an inconsistent state"
            echo -e "${RED}${BOLD}[WARNING]${NC} Please check your configuration manually"
        else
            echo -e "${GREEN}[SUCCESS]${NC} Successfully restored previous configuration"
        fi
        exit 1
    fi

    if ! iptables -t nat -A POSTROUTING -s "${target_ip}" -j SNAT --to-source "${source_ip}"; then
        echo -e "${RED}${BOLD}[ERROR]${NC} Failed to set up SNAT rules"
        echo -e "${YELLOW}[INFO]${NC} Attempting to restore previous configuration..."
        if ! iptables-restore < "$backup_file"; then
            echo -e "${RED}${BOLD}[ERROR]${NC} Failed to restore from backup"
            exit 1
        fi
        exit 1
    fi

    # Set up FORWARD chain
    echo -e "${CYAN}Setting up forwarding rules...${NC}"
    iptables -A FORWARD -i ${source_interface} -o tailscale0 -j ACCEPT
    iptables -A FORWARD -i tailscale0 -o ${source_interface} -j ACCEPT
    iptables -A FORWARD -p icmp -j ACCEPT
    iptables -A FORWARD -m state --state RELATED,ESTABLISHED -j ACCEPT

    # Add Tailscale subnet handling
    echo -e "${CYAN}Setting up Tailscale subnet rules...${NC}"
    iptables -t nat -A POSTROUTING -s 100.64.0.0/10 -j SNAT --to-source "${source_ip}"
    iptables -t nat -A POSTROUTING -s "${source_subnet}" -j MASQUERADE

    # Verify configuration
    echo -e "${YELLOW}[4/7]${NC} Verifying configuration..."
    if ! iptables -t nat -L PREROUTING -n -v >/dev/null 2>&1; then
        echo -e "${RED}${BOLD}[ERROR]${NC} Configuration verification failed. Restoring previous rules..."
        if ! iptables-restore < "$backup_file"; then
            echo -e "${RED}${BOLD}[ERROR]${NC} Failed to restore from backup"
            exit 1
        fi
        exit 1
    fi

    # Enable service
    echo -e "${YELLOW}[5/7]${NC} Enabling iptables-persistent service..."
    if systemctl is-active --quiet netfilter-persistent; then
        systemctl restart netfilter-persistent
    else
        systemctl enable netfilter-persistent
        systemctl start netfilter-persistent
    fi

    # Save current configuration
    echo -e "${YELLOW}[6/7]${NC} Saving current configuration..."
    if ! save_iptables; then
        echo -e "${RED}${BOLD}[ERROR]${NC} Failed to save configuration. Current rules may be lost on reboot"
        exit 1
    fi

    # Final status
    echo -e "${YELLOW}[7/7]${NC} Configuration complete"
    echo -e "\n${GREEN}${BOLD}[SUCCESS]${NC} Configuration completed successfully!"
    echo -e "${CYAN}Preserved TCP ports: ${PRESERVED_TCP_PORTS[*]}${NC}"
    echo -e "${CYAN}Preserved UDP ports: ${PRESERVED_UDP_PORTS[*]}${NC}"
    echo -e "${CYAN}All other ports are forwarded to the Tailscale device${NC}"
    echo -e "${CYAN}Backup saved to: $backup_file${NC}\n"
else
    echo -e "\n${YELLOW}Operation cancelled by user${NC}"
    exit 1
fi