#!/bin/bash
# Generate WireGuard vanity keys and update the vault file
# Usage: ./generate-wireguard-keys.sh

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_DIR="$(dirname "$SCRIPT_DIR")"
VAULT_FILE="$REPO_DIR/vault/wireguard.yml"
TEMP_FILE=$(mktemp)

echo "==================================="
echo "WireGuard Vanity Key Generator"
echo "==================================="
echo ""

# Check if wireguard-vanity-address is installed
if ! command -v wireguard-vanity-address &> /dev/null; then
    echo "Error: wireguard-vanity-address not found"
    echo "Install it with: cargo install wireguard-vanity-address"
    exit 1
fi

# Check if wg is installed
if ! command -v wg &> /dev/null; then
    echo "Error: wg (wireguard-tools) not found"
    echo "Install it with: sudo apt install wireguard-tools"
    exit 1
fi

# Function to generate a vanity keypair
generate_vanity_key() {
    local name=$1
    local search=$2
    local search_len=${#search}
    local timeout=${3:-120}

    echo "Generating vanity key for $name (searching for '$search')..."

    # Try to generate vanity key with timeout
    local output=$(timeout $timeout wireguard-vanity-address --in $search_len "$search" 2>&1 | grep "^private" | head -n 1 || true)

    if [ -n "$output" ]; then
        local private_key=$(echo "$output" | awk '{print $2}')
        local public_key=$(echo "$output" | awk '{print $4}')
        echo "  ✓ Found vanity key: $public_key"
    else
        echo "  ⚠ No vanity key found in ${timeout}s, generating regular key..."
        local private_key=$(wg genkey)
        local public_key=$(echo "$private_key" | wg pubkey)
        echo "  ✓ Generated regular key: $public_key"
    fi

    echo "$private_key|$public_key"
}

echo "This will generate new WireGuard keys. This will:"
echo "  1. Generate vanity keys for all peers"
echo "  2. Update vault/wireguard.yml with the new keys"
echo ""
read -p "Continue? (y/N) " -n 1 -r
echo
if [[ ! $REPLY =~ ^[Yy]$ ]]; then
    echo "Aborted."
    exit 0
fi
echo ""

# Generate keys
echo "Generating keys (this may take a few minutes)..."
echo ""

SERVER_KEYS=$(generate_vanity_key "server" "talos" 180)
SERVER_PRIVATE=$(echo "$SERVER_KEYS" | cut -d'|' -f1)
SERVER_PUBLIC=$(echo "$SERVER_KEYS" | cut -d'|' -f2)
echo ""

CLIENT1_KEYS=$(generate_vanity_key "client1" "lptp" 120)
CLIENT1_PRIVATE=$(echo "$CLIENT1_KEYS" | cut -d'|' -f1)
CLIENT1_PUBLIC=$(echo "$CLIENT1_KEYS" | cut -d'|' -f2)
echo ""

CLIENT2_KEYS=$(generate_vanity_key "client2" "phon" 120)
CLIENT2_PRIVATE=$(echo "$CLIENT2_KEYS" | cut -d'|' -f1)
CLIENT2_PUBLIC=$(echo "$CLIENT2_KEYS" | cut -d'|' -f2)
echo ""

MASSER_KEYS=$(generate_vanity_key "masser" "mssr" 120)
MASSER_PRIVATE=$(echo "$MASSER_KEYS" | cut -d'|' -f1)
MASSER_PUBLIC=$(echo "$MASSER_KEYS" | cut -d'|' -f2)
echo ""

# Create the vault file content
cat > "$TEMP_FILE" << EOF
---
# WireGuard VPN Keys
# These are encrypted with ansible-vault
# To edit: ansible-vault edit vault/wireguard.yml

wireguard_server:
  private_key: "$SERVER_PRIVATE"
  public_key: "$SERVER_PUBLIC"
  vanity_search: "talos"

wireguard_clients:
  - name: client1
    private_key: "$CLIENT1_PRIVATE"
    public_key: "$CLIENT1_PUBLIC"
    vanity_search: "lptp"
    ip_address: "10.8.0.2/24"
    description: "Example laptop client"

  - name: client2
    private_key: "$CLIENT2_PRIVATE"
    public_key: "$CLIENT2_PUBLIC"
    vanity_search: "phon"
    ip_address: "10.8.0.3/24"
    description: "Example phone client"

  - name: masser
    private_key: "$MASSER_PRIVATE"
    public_key: "$MASSER_PUBLIC"
    vanity_search: "mssr"
    ip_address: "10.8.0.4/24"
    description: "Masser phone"
EOF

echo "==================================="
echo "Keys generated successfully!"
echo "==================================="
echo ""
echo "Summary:"
echo "  Server:  $SERVER_PUBLIC"
echo "  Client1: $CLIENT1_PUBLIC"
echo "  Client2: $CLIENT2_PUBLIC"
echo "  Masser:  $MASSER_PUBLIC"
echo ""

# Check if vault file should be encrypted
if [ -f "$VAULT_FILE" ]; then
    # Check if existing file is encrypted
    if head -n 1 "$VAULT_FILE" | grep -q '^\$ANSIBLE_VAULT'; then
        echo "Encrypting with ansible-vault..."
        ansible-vault encrypt "$TEMP_FILE" --output="$VAULT_FILE"
        rm -f "$TEMP_FILE"
        echo "✓ Keys saved to $VAULT_FILE (encrypted)"
    else
        echo "⚠ Warning: Existing vault file is not encrypted"
        echo "Saving unencrypted keys to $VAULT_FILE"
        mv "$TEMP_FILE" "$VAULT_FILE"
        echo ""
        echo "To encrypt the file, run:"
        echo "  ansible-vault encrypt $VAULT_FILE"
    fi
else
    echo "Do you want to encrypt the vault file with ansible-vault? (recommended)"
    read -p "Encrypt? (Y/n) " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Nn]$ ]]; then
        ansible-vault encrypt "$TEMP_FILE" --output="$VAULT_FILE"
        rm -f "$TEMP_FILE"
        echo "✓ Keys saved to $VAULT_FILE (encrypted)"
    else
        mv "$TEMP_FILE" "$VAULT_FILE"
        echo "✓ Keys saved to $VAULT_FILE (unencrypted)"
        echo ""
        echo "⚠ Warning: Keys are stored unencrypted!"
        echo "To encrypt them later, run:"
        echo "  ansible-vault encrypt $VAULT_FILE"
    fi
fi

echo ""
echo "Next steps:"
echo "  1. Run the playbook to deploy: ansible-playbook -i inventory.yml playbook.yml --limit talos --tags leaf"
echo "  2. Update your WireGuard clients with the new server public key"
echo ""