#!/bin/bash

################################################################################
# File Name    : wireguard-namespace-qube.sh
# Description  : WireGuard VPN setup for Qubes OS that routes all internet traffic
#                through the wg0 interface using network namespace isolation to
#                prevent DNS and IP leaks. Configures WireGuard directly via ip
#                commands instead of wg-quick, implements minimal NFTables rules
#                to block non-VPN traffic and enable DNS NAT, with IPv6 disabled.
# Usage        : • Transfer this script from appvm to dom0 with:
#                [user@dom0 ~]$ qvm-run -p appvm 'cat ~/wireguard-namespace-qube.sh' > ~/wireguard-namespace-qube.sh
#                • Place your WireGuard config to the desired folder
#                [user@dom0 ~]$ mkdir -p wireguard
#                [user@dom0 ~]$ qvm-run -p appvm 'cat ~/wg0.conf' > ~/wireguard/wg0.conf
#                • Make the script executable with:
#                [user@dom0 ~]$ chmod +x ~/wireguard-namespace-qube.sh
#                • Run the script with:
#                [user@dom0 ~]$ bash ~/wireguard-namespace-qube.sh
# Author       : Me and the bois
# License      : Free of charge, no warranty
# Last edited  : 2025-10-18
################################################################################

# Safety first
set -euo pipefail

# Configuration (set default values here)
WG_CONFIG_NAME="wg0.conf"
WG_CONFIG_SOURCE="/home/user/wireguard"
BASE_TEMPLATE="debian-12-minimal"
CUSTOM_TEMPLATE="debian-vpn-template"
VPN_SERVICE="sys-vpn"

# Step 1: Verify and create base template
create_custom_template() {
    echo -e "\nStep 1: Verifying and creating: $BASE_TEMPLATE..."

    if ! qvm-check "$BASE_TEMPLATE"; then
        echo -e "\nInstalling $BASE_TEMPLATE..."
        sudo qubes-dom0-update "qubes-template-$BASE_TEMPLATE"
    fi

    qvm-shutdown --wait "$BASE_TEMPLATE" 2>/dev/null || true

    if qvm-check "$BASE_TEMPLATE"; then
        echo -e "\nUpdating $BASE_TEMPLATE..."
        sudo qubesctl --show-output --skip-dom0 --targets="$BASE_TEMPLATE" state.sls update.qubes-vm
    fi

    echo -e "\nSetup complete"
}

# Step 2: Verify and create custom template with packages (change as needed)
install_packages() {
    echo -e "\nStep 2: Creating and installing required packages in $CUSTOM_TEMPLATE..."

    if qvm-check "$CUSTOM_TEMPLATE"; then
        echo -e "\n$CUSTOM_TEMPLATE already exists"
        echo -e "\nContinuing to $VPN_SERVICE creation..."
    else
        echo -e "\nCreating $CUSTOM_TEMPLATE by cloning from $BASE_TEMPLATE..."
        qvm-clone "$BASE_TEMPLATE" "$CUSTOM_TEMPLATE"
        qvm-prefs "$CUSTOM_TEMPLATE" label orange

        echo -e "\nStarting $CUSTOM_TEMPLATE for package installation..."
        qvm-start "$CUSTOM_TEMPLATE"

        echo -e "\nConfiguring template..."
        qvm-run -p -u root "$CUSTOM_TEMPLATE" "echo 'TERM=xterm' >> /etc/environment"
        qvm-run -p -u root "$CUSTOM_TEMPLATE" "sed -i 's/^# *\(en_US.UTF-8\)/\1/' /etc/locale.gen"
        qvm-run -p -u root "$CUSTOM_TEMPLATE" "locale-gen en_US.UTF-8"

        echo -e "\nInstalling packages in $CUSTOM_TEMPLATE..."
        qvm-run -p -u root "$CUSTOM_TEMPLATE" "apt update -y"
        qvm-run -p -u root "$CUSTOM_TEMPLATE" "apt install -y --no-install-recommends \
            qubes-core-agent-passwordless-root \
            qubes-core-agent-networking \
            wireguard-tools \
            nftables \
            iproute2 \
            dnsutils \
            traceroute \
            less \
            psmisc \
            xfce4-terminal"

        echo -e "\n$CUSTOM_TEMPLATE created successfully"
    fi

    qvm-shutdown --wait "$CUSTOM_TEMPLATE" 2>/dev/null || true
}

# Step 3: Create qube service
create_vpn_qube() {
    echo -e "\nStep 3: Verifying and creating qube service: $VPN_SERVICE..."

    if qvm-check "$VPN_SERVICE"; then
        echo -e "\n$VPN_SERVICE already exists"
        echo -e "\nExiting"
        exit 1
    else
        echo -e "\nCreating $VPN_SERVICE..."

        qvm-create \
          --class AppVM \
          --template "$CUSTOM_TEMPLATE" \
          --label orange \
          --property netvm=sys-firewall \
          --property provides_network=True \
          --property vcpus=1 \
          --property memory=300 \
          --property maxmem=400 \
          --property autostart=False \
          --property include_in_backups=False \
          "$VPN_SERVICE"

        qvm-features "$VPN_SERVICE" ipv6 ''
        echo -e "\nDisabled IPv6 in $VPN_SERVICE"

        qvm-service "$VPN_SERVICE" qubes-firewall on

        echo -e "\nStarting $VPN_SERVICE to apply configuration..."
        qvm-start "$VPN_SERVICE"

        echo -e "\n$VPN_SERVICE created and configured"
    fi
}

# Step 4: Comment out to configure via ip commands instead of wg-quick
comment_wg_config() {
    echo -e "\nStep 4: Commenting out Address, DNS, and MTU in WireGuard config..."

    local temp_config=$(mktemp)
    cp "$WG_CONFIG_SOURCE/$WG_CONFIG_NAME" "$temp_config"

    # Comment out the lines
    sed -i 's/^Address/#Address/g' "$temp_config"
    sed -i 's/^DNS/#DNS/g' "$temp_config"
    sed -i 's/^MTU/#MTU/g' "$temp_config"

    # Overwrite the original
    mv "$temp_config" "$WG_CONFIG_SOURCE/$WG_CONFIG_NAME"
    echo "WireGuard config automatically commented"
}

# Step 5: Copy the configuration file from dom0 to the VPN service qube
transfer_config() {
    echo -e "\nStep 5: Transferring WireGuard configuration to the VPN service qube..."

    if [[ ! -f "$WG_CONFIG_SOURCE/$WG_CONFIG_NAME" ]]; then
        echo -e "\nError: Config file not found at $WG_CONFIG_SOURCE/$WG_CONFIG_NAME"
        return 1
    fi

    echo -e "\nCopying..."
    qvm-copy-to-vm "$VPN_SERVICE" "$WG_CONFIG_SOURCE/$WG_CONFIG_NAME"
    sleep 2

    qvm-run -p -u root "$VPN_SERVICE" "mkdir -p /rw/config"
    qvm-run -p -u root "$VPN_SERVICE" "mkdir -p /rw/config/vpn"
    qvm-run -p -u root "$VPN_SERVICE" "mv '/home/user/QubesIncoming/dom0/$WG_CONFIG_NAME' '/rw/config/vpn/'"
    qvm-run -p -u root "$VPN_SERVICE" "chmod 600 '/rw/config/vpn/$WG_CONFIG_NAME'"
    echo -e "\nConfig file moved to: /rw/config/vpn/$WG_CONFIG_NAME"
}

# Create WireGuard namespace setup script
create_wireguard_setup_script() {
    echo -e "\nCreating WireGuard namespace setup script..."

    local temp_script=$(mktemp)
    cat > "$temp_script" << 'EOF'
#!/bin/bash
set -euo pipefail

echo -e "\nStarting WireGuard namespace setup"
sleep 5

# The WG_CONFIG_NAME variable should be passed via export from the main script
if [ -z "$WG_CONFIG_NAME" ]; then
    echo "Error: WG_CONFIG_NAME environment variable not set"
    exit 1
fi

INTERFACE="eth0"
echo -e "\nUsing interface: $INTERFACE"

# Get network configuration
PHYSICAL_IP=$(ip -4 addr show "$INTERFACE" | grep -oP "(?<=inet\s)\d+(\.\d+){3}/\d+")
PHYSICAL_ROUTE=$(ip route | grep default | awk '{print $3}')
WG_CONFIG="/rw/config/vpn/$WG_CONFIG_NAME"

# Extract WireGuard IP
WG_IP=$(grep -E "^(#)?Address" "$WG_CONFIG" | head -n1 | cut -d"=" -f2 | cut -d"/" -f1 | tr -d " ")
if [ -z "$WG_IP" ]; then
    echo "Error: No Address found in WireGuard config"
    exit 1
fi
WG_IP="${WG_IP}/32"
echo -e "\nUsing WG_IP: $WG_IP"

# Extract DNS for endpoint resolution
DNS_LINE=$(grep -E "^(#\s*)?DNS" "$WG_CONFIG" | head -1)
if [ -n "$DNS_LINE" ]; then
    CLEAN_DNS_LINE=$(echo "$DNS_LINE" | sed -E "s/^#\s*//")
    DNS_SERVERS=$(echo "$CLEAN_DNS_LINE" | cut -d= -f2 | sed "s/^[[:space:]]*//;s/[[:space:]]*\$//" | tr "," " ")
    WG_DNS=$(echo "$DNS_SERVERS" | cut -d" " -f1)
    echo -e "\nUsing DNS from WireGuard config for resolution: $DNS_SERVERS"
else
    WG_DNS="9.9.9.9"
    echo -e "\nUsing fallback DNS for resolution: $WG_DNS"
fi

# Extract endpoint
ENDPOINT=$(grep "^Endpoint" "$WG_CONFIG" | cut -d"=" -f2 | tr -d " " | cut -d: -f1)
PORT=$(grep "^Endpoint" "$WG_CONFIG" | cut -d"=" -f2 | tr -d " " | cut -d: -f2)

RESOLVED_IP=""
if [ -n "$ENDPOINT" ]; then
    echo -e "\nResolving endpoint: $ENDPOINT"
    for dns_server in $WG_DNS 10.139.1.1 9.9.9.9; do
        IP=$(dig +short +time=2 +tries=2 "$ENDPOINT" "@$dns_server" 2>/dev/null | grep -E "^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+\$" | head -1)
        if [ -n "$IP" ] && [ "$IP" != "127.0.0.1" ]; then
            echo -e "\nResolved $ENDPOINT to $IP using DNS $dns_server"
            RESOLVED_IP="$IP"
            break
        fi
    done
fi

if [ -z "$RESOLVED_IP" ]; then
    echo -e "\nError: Could not resolve endpoint $ENDPOINT"
    exit 1
fi

echo -e "\nSetting up network namespace..."

# Clean up any existing namespace
ip netns delete physical 2>/dev/null || true

# Create namespace and move interface
ip netns add physical
ip link set "$INTERFACE" netns physical
ip -n physical link add wg0 type wireguard
ip -n physical link set wg0 netns 1

# Reconfigure physical interface in namespace
ip -n physical addr add "$PHYSICAL_IP" dev "$INTERFACE"
ip -n physical link set "$INTERFACE" up
ip -n physical route add default via "$PHYSICAL_ROUTE" dev "$INTERFACE" onlink

# Create temporary config with resolved IP
cp "$WG_CONFIG" /tmp/wg-config-temp.conf
sed -i "s/$ENDPOINT:[0-9]*/$RESOLVED_IP:$PORT/g" /tmp/wg-config-temp.conf

# Configure WireGuard (remove conflicting lines)
sed -i "/^Address/d;/^DNS/d;/^MTU/d" /tmp/wg-config-temp.conf
wg setconf wg0 /tmp/wg-config-temp.conf
rm -f /tmp/wg-config-temp.conf

ip addr add "$WG_IP" dev wg0
ip link set wg0 up

# Set up routing - all traffic through WireGuard
ip route del default 2>/dev/null || true
ip route add default dev wg0

echo -e "\nWireGuard namespace setup complete - all traffic routed through VPN"
EOF

    qvm-copy-to-vm "$VPN_SERVICE" "$temp_script"
    qvm-run -p -u root "$VPN_SERVICE" "mv '/home/user/QubesIncoming/dom0/$(basename $temp_script)' '/rw/config/setup-wireguard-namespace.sh'"
    qvm-run -p -u root "$VPN_SERVICE" "chown root:root '/rw/config/setup-wireguard-namespace.sh'"
    qvm-run -p -u root "$VPN_SERVICE" "chmod 755 '/rw/config/setup-wireguard-namespace.sh'"
    rm -f "$temp_script"
    echo "WireGuard namespace setup script created successfully"
}

# Create NFTables configuration (variable expansion with mktemp strategy) (using echo commands to avoid heredoc issues and complexity)
create_nftables_generator_script() {
    echo -e "\nCreating NFTables generator script..."

    local temp_script=$(mktemp)
    cat > "$temp_script" << 'EOF'
#!/bin/bash
set -e

echo "Creating NFTables configuration..."

# The WG_CONFIG_NAME variable should be passed via export from the main script
if [ -z "$WG_CONFIG_NAME" ]; then
    echo "Error: WG_CONFIG_NAME environment variable not set"
    exit 1
fi

# Extract DNS from WireGuard config
WG_CONFIG="/rw/config/vpn/$WG_CONFIG_NAME"

# Check if config file exists
if [ ! -f "$WG_CONFIG" ]; then
    echo "Error: WireGuard config file not found at $WG_CONFIG"
    exit 1
fi

# Extract DNS servers
DNS_LINE=$(grep -E "^(#\s*)?DNS" "$WG_CONFIG" | head -1)
if [ -n "$DNS_LINE" ]; then
    CLEAN_DNS_LINE=$(echo "$DNS_LINE" | sed -E "s/^#\s*//")
    DNS_SERVERS=$(echo "$CLEAN_DNS_LINE" | cut -d'=' -f2 | sed 's/^[[:space:]]*//;s/[[:space:]]*$//' | tr ',' ' ')
    VPNDNS1=$(echo "$DNS_SERVERS" | awk '{print $1}')
    VPNDNS2=$(echo "$DNS_SERVERS" | awk '{print $2}')
else
    echo "Error: No DNS found in WireGuard config"
    exit 1
fi

echo "Using DNS from WireGuard config: $DNS_SERVERS"
echo "Primary DNS: $VPNDNS1"
echo "Secondary DNS: $VPNDNS2"

# Create nftables.conf
{
echo '#!/usr/sbin/nft -f'
echo ''
echo 'flush chain ip qubes dnat-dns'
echo ''
echo 'table ip qubes {'
echo '    chain dnat-dns {'
echo '        type nat hook prerouting priority dstnat; policy accept;'
echo "        ip daddr 10.139.1.1 udp dport 53 dnat to $VPNDNS1"
echo "        ip daddr 10.139.1.1 tcp dport 53 dnat to $VPNDNS1"
if [[ -n "$VPNDNS2" && "$VPNDNS2" =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
echo "        ip daddr 10.139.1.2 udp dport 53 dnat to $VPNDNS2"
echo "        ip daddr 10.139.1.2 tcp dport 53 dnat to $VPNDNS2"
fi
echo '    }'
echo '}'
} > /rw/config/nftables.conf

echo "NFTables configuration created successfully"
EOF

    qvm-copy-to-vm "$VPN_SERVICE" "$temp_script"
    qvm-run -p -u root "$VPN_SERVICE" "mv '/home/user/QubesIncoming/dom0/$(basename $temp_script)' '/rw/config/nftables-generator.sh'"
    qvm-run -p -u root "$VPN_SERVICE" "chown root:root '/rw/config/nftables-generator.sh'"
    qvm-run -p -u root "$VPN_SERVICE" "chmod 755 '/rw/config/nftables-generator.sh'"
    rm -f "$temp_script"
    echo "NFTables generator script created successfully"
}

# Create nftables rc autostart script
create_nftables_rc_autostart() {
    echo -e "\nCreating NFTables RC autostart script..."

    local temp_script=$(mktemp)
    cat > "$temp_script" << EOF
#!/bin/bash

echo -e "\nStarting WireGuard VPN setup..."
sleep 2

if ip link show wg0 >/dev/null 2>&1; then
    echo -e "\nWireGuard is already running"
    exit 0
fi

echo -e "\nRunning WireGuard setup..."
export WG_CONFIG_NAME="$WG_CONFIG_NAME"
/rw/config/setup-wireguard-namespace.sh
echo -e "\nWireGuard setup completed"
sleep 5
EOF

    qvm-run -p -u root "$VPN_SERVICE" "mkdir -p /rw/config/rc.local.d"
    qvm-copy-to-vm "$VPN_SERVICE" "$temp_script"
    qvm-run -p -u root "$VPN_SERVICE" "mv '/home/user/QubesIncoming/dom0/$(basename $temp_script)' '/rw/config/rc.local.d/nftables-rc-autostart.rc'"
    qvm-run -p -u root "$VPN_SERVICE" "chown root:root '/rw/config/rc.local.d/nftables-rc-autostart.rc'"
    qvm-run -p -u root "$VPN_SERVICE" "chmod 755 '/rw/config/rc.local.d/nftables-rc-autostart.rc'"

    rm -f "$temp_script"
    echo "NFTables RC autostart script created successfully"
}

# Step 6: Finalize with DNS test
finalize() {
    echo -e "\nStep 6: Finalizing setup..."

    # Clean up
    qvm-run -p -u root "$VPN_SERVICE" "rmdir '/home/user/QubesIncoming/dom0' 2>/dev/null || true"
    qvm-run -p -u root "$VPN_SERVICE" "rmdir '/home/user/QubesIncoming' 2>/dev/null || true"

    # Generate NFTables config (using export to make the variable available to the script)
    echo -e "\nGenerating NFTables DNS configuration..."
    qvm-run -p -u root "$VPN_SERVICE" "export WG_CONFIG_NAME='$WG_CONFIG_NAME'; /rw/config/nftables-generator.sh"

    # Apply NFTables config
    echo -e "\nApplying DNS redirection rules..."
    qvm-run -p -u root "$VPN_SERVICE" "nft -f /rw/config/nftables.conf"

    # Verify the rules were applied
    echo -e "\nVerifying NFTables rules..."
    qvm-run -p -u root "$VPN_SERVICE" "nft list chain ip qubes dnat-dns"

    # Setup WireGuard (using export to make the variable available to the script)
    echo -e "\nSetting up WireGuard namespace..."
    qvm-run -p -u root "$VPN_SERVICE" "export WG_CONFIG_NAME='$WG_CONFIG_NAME'; /rw/config/setup-wireguard-namespace.sh"

    # Test connection
    echo -e "\nTesting connection..."
    sleep 5
    qvm-run -p "$VPN_SERVICE" "ping -c 9.9.9.9"
    echo -e "\nTesting connection with MTU 1420..."
    qvm-run -p "$VPN_SERVICE" "ping -M do -s 1420 9.9.9.9"
}

# Main function
main() {
    echo -e "\nStarting WireGuard VPN setup..."

    create_custom_template
    install_packages
    create_vpn_qube
    comment_wg_config
    transfer_config
    create_wireguard_setup_script
    create_nftables_generator_script
    create_nftables_rc_autostart
    finalize

    echo -e "\nSetup completed successfully!"
}

# Run
main "$@"