#!/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 IP and DNS leaks. Below set the desired
#                values for your variables, note that you can rerun the
#                script to create other VPN services.
# 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-28
################################################################################

# Safety first
set -euo pipefail

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

# 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 install -y --no-install-recommends \
      qubes-core-agent-passwordless-root \
      qubes-core-agent-networking \
      wireguard-tools \
      nftables \
      iproute2 \
      dnsutils \
      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"

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

    echo -e "\nEnabling qubes-firewall service 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 to configure via ip commands"
}

# 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 must be passed via export from the main script
if [ -z "$WG_CONFIG_NAME" ]; then
    echo "Error: WG_CONFIG_NAME environment variable not set"
    echo "This script must be called with: export WG_CONFIG_NAME='wg0.conf'; /rw/config/setup-wireguard-namespace.sh"
    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 generator script that uses exported variables
create_nftables_generator_script() {
    echo -e "\nCreating NFTables generator script..."

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

echo "Generating NFTables configuration..."

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

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

# Get physical IP for virtualif (same as PHYSICAL_IP from setup script)
VIRTUALIF=$(ip -4 addr show eth0 | grep -oP "(?<=inet\s)\d+(\.\d+){3}/\d+" | head -1)
if [ -z "$VIRTUALIF" ]; then
    echo "Error: Could not determine physical IP"
    exit 1
fi

# Extract DNS servers from WireGuard config
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*//")
    VPNDNS1=$(echo "$CLEAN_DNS_LINE" | cut -d'=' -f2 | sed 's/^[[:space:]]*//;s/[[:space:]]*$//' | tr ',' ' ' | awk '{print $1}')
else
    echo "Error: No DNS found in WireGuard config"
    exit 1
fi

echo "Using physical IP: $VIRTUALIF"
echo "Using VPN DNS: $VPNDNS1"

# Generate nftables config file
cat > /rw/config/nftables.conf << NFTEOF
#!/usr/sbin/nft -f

# Set variables
define virtualif = $VIRTUALIF
define vpndns1 = $VPNDNS1

# Flush existing rules in qubes dnat-dns chain
flush chain qubes dnat-dns

# Table for NAT and vpn-dns-redirect chain
table ip nat {
    chain vpn-dns-redirect {
        type nat hook prerouting priority -100; policy accept;
        ip daddr 10.139.1.1 udp dport 53 dnat to \$vpndns1
        ip daddr 10.139.1.1 tcp dport 53 dnat to \$vpndns1
    }
}

# Mangle table (for MSS clamping)
table ip mangle {
    chain forward {
        type filter hook forward priority 0; policy accept;
        tcp flags syn tcp option maxseg size set rt mtu;
    }
}
NFTEOF

echo "NFTables configuration generated successfully at /rw/config/nftables.conf"
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 main startup script that uses exported variables
create_main_startup_script() {
    echo -e "\nCreating main startup script..."

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

echo "Starting WireGuard VPN setup..."

# Check if WG_CONFIG_NAME is set in environment
if [ -z "$WG_CONFIG_NAME" ]; then
    echo "Error: WG_CONFIG_NAME environment variable not set"
    echo "This script must be called with: export WG_CONFIG_NAME='wg0.conf'; /rw/config/start-wireguard.sh"
    exit 1
fi

echo "Using WireGuard config: $WG_CONFIG_NAME"

# Generate NFTables configuration
echo "Generating NFTables DNS configuration..."
/rw/config/nftables-generator.sh

# Apply NFTables config
echo "Applying DNS redirection rules..."
nft -f /rw/config/nftables.conf

# Setup WireGuard namespace
echo "Setting up WireGuard namespace..."
/rw/config/setup-wireguard-namespace.sh

echo "WireGuard VPN setup completed"
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/start-wireguard.sh'"
    qvm-run -p -u root "$VPN_SERVICE" "chown root:root '/rw/config/start-wireguard.sh'"
    qvm-run -p -u root "$VPN_SERVICE" "chmod 755 '/rw/config/start-wireguard.sh'"
    rm -f "$temp_script"
    echo "Main startup script completed"
}

# Create rc.local entry that sets the environment variable
create_rclocal_entry() {
    echo -e "\nCreating rc.local entry..."

    local temp_script=$(mktemp)
    cat > "$temp_script" << EOF
#!/bin/bash
# Start WireGuard VPN on boot
export WG_CONFIG_NAME="$WG_CONFIG_NAME"
/rw/config/start-wireguard.sh
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/rc.local'"
    qvm-run -p -u root "$VPN_SERVICE" "chown root:root '/rw/config/rc.local'"
    qvm-run -p -u root "$VPN_SERVICE" "chmod 755 '/rw/config/rc.local'"
    rm -f "$temp_script"
    echo "rc.local entry created successfully"
}

# Finalize
finalize() {
    echo -e "\nFinalizing..."

    # 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 nat vpn-dns-redirect"

    # 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"

    # Verify WireGuard is running
    echo -e "\nVerifying WireGuard connection..."
    qvm-run -p -u root "$VPN_SERVICE" "wg show"

    # Configure menu items
    qvm-features "$CUSTOM_TEMPLATE" menu-items "xfce4-terminal.desktop"
    qvm-features "$VPN_SERVICE" menu-items "xfce4-terminal.desktop"

    # Shutdown
    qvm-shutdown --wait "$VPN_SERVICE" 2>/dev/null || true
}

# 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_main_startup_script
    create_rclocal_entry
    finalize

    echo -e "\nSetup completed!"
}

# Run
main "$@"