tunnel-setup-relay.sh

#!/usr/bin/env bash
# tunnel-setup-relay.sh
# Configure a WireGuard relay server.
#
# Supports multiple provisioning modes:
#   --provider=digitalocean  Provision a new DigitalOcean Droplet (requires DO_API_TOKEN)
#   --provider=ssh           Configure an existing server via SSH (requires RELAY_HOST)
#   --provider=local         Configure WireGuard on this machine (the relay IS this machine)
#   --generate-script        Output the setup script without executing (for manual use)
#
# Default: digitalocean (if DO_API_TOKEN is set), otherwise prompts for mode.
#
# Copyright 2026 Quantum Pipes Technologies, LLC
# SPDX-License-Identifier: Apache-2.0

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# shellcheck disable=SC1091
source "$SCRIPT_DIR/tunnel-preflight.sh"

# ---------------------------------------------------------------------------
# Parse arguments
# ---------------------------------------------------------------------------
PROVIDER=""
RELAY_HOST="${RELAY_HOST:-}"
RELAY_SSH_USER="${RELAY_SSH_USER:-root}"
GENERATE_SCRIPT=false

for arg in "$@"; do
    case "$arg" in
        --provider=*) PROVIDER="${arg#*=}" ;;
        --host=*) RELAY_HOST="${arg#*=}" ;;
        --user=*) RELAY_SSH_USER="${arg#*=}" ;;
        --generate-script) GENERATE_SCRIPT=true ;;
        --help|-h)
            echo "Usage: tunnel-setup-relay.sh [OPTIONS]"
            echo ""
            echo "Options:"
            echo "  --provider=MODE     Provisioning mode: digitalocean, ssh, local"
            echo "  --host=IP           Server IP/hostname (required for --provider=ssh)"
            echo "  --user=USER         SSH user (default: root)"
            echo "  --generate-script   Output setup script to stdout (no execution)"
            echo ""
            echo "Environment:"
            echo "  DO_API_TOKEN        DigitalOcean API token (for --provider=digitalocean)"
            echo "  RELAY_HOST          Server IP (alternative to --host=)"
            echo "  RELAY_SSH_USER      SSH user (alternative to --user=)"
            exit 0
            ;;
    esac
done

# Auto-detect provider if not specified
if [[ -z "$PROVIDER" ]]; then
    if [[ "$GENERATE_SCRIPT" == "true" ]]; then
        PROVIDER="script"
    elif [[ -n "${DO_API_TOKEN:-}" ]]; then
        PROVIDER="digitalocean"
    elif [[ -n "$RELAY_HOST" ]]; then
        PROVIDER="ssh"
    else
        log_error "No provider specified. Use --provider=digitalocean, --provider=ssh, or --provider=local"
        log_error "  DigitalOcean: export DO_API_TOKEN=... && tunnel-setup-relay.sh"
        log_error "  Existing VPS: tunnel-setup-relay.sh --provider=ssh --host=YOUR_SERVER_IP"
        log_error "  This machine: tunnel-setup-relay.sh --provider=local"
        log_error "  Script only:  tunnel-setup-relay.sh --generate-script"
        exit 1
    fi
fi

# ---------------------------------------------------------------------------
# Generate relay keypair (shared across all providers)
# ---------------------------------------------------------------------------
set_safe_umask
keypair="$(wg_genkeypair)"
RELAY_PRIVATE_KEY="${keypair%% *}"
RELAY_PUBLIC_KEY="${keypair##* }"

config_dir="$(ensure_config_dir)"
printf '%s\n' "$RELAY_PRIVATE_KEY" > "${config_dir}/relay.key"
chmod 600 "${config_dir}/relay.key"
printf '%s\n' "$RELAY_PUBLIC_KEY" > "${config_dir}/relay.pub"

log_info "Relay keypair generated"

# ---------------------------------------------------------------------------
# Build the setup script (used by all providers)
# ---------------------------------------------------------------------------
_build_setup_script() {
    cat <<SETUP_SCRIPT
#!/bin/bash
# QP Tunnel Relay Setup Script
# Generated by tunnel-setup-relay.sh
# Run this on your relay server as root.
set -euo pipefail

echo "[1/6] Installing WireGuard..."
if command -v apt-get &>/dev/null; then
    apt-get update -y
    apt-get install -y wireguard qrencode
elif command -v dnf &>/dev/null; then
    dnf install -y wireguard-tools qrencode
elif command -v yum &>/dev/null; then
    yum install -y wireguard-tools qrencode
elif command -v pacman &>/dev/null; then
    pacman -Sy --noconfirm wireguard-tools qrencode
elif command -v apk &>/dev/null; then
    apk add wireguard-tools qrencode
else
    echo "ERROR: Unsupported package manager. Install wireguard-tools manually."
    exit 1
fi

echo "[2/6] Writing WireGuard config..."
umask 077
cat > /etc/wireguard/${TUNNEL_INTERFACE}.conf <<'WGCONF'
[Interface]
PrivateKey = ${RELAY_PRIVATE_KEY}
Address = ${TUNNEL_RELAY_IP}/24
ListenPort = ${TUNNEL_PORT}
SaveConfig = true
PostUp = iptables -A FORWARD -i ${TUNNEL_INTERFACE} -j ACCEPT; iptables -t nat -A POSTROUTING -o \$(ip route show default | awk '/default/ {print \$5}' | head -1) -j MASQUERADE
PostDown = iptables -D FORWARD -i ${TUNNEL_INTERFACE} -j ACCEPT; iptables -t nat -D POSTROUTING -o \$(ip route show default | awk '/default/ {print \$5}' | head -1) -j MASQUERADE
WGCONF
chmod 600 /etc/wireguard/${TUNNEL_INTERFACE}.conf

echo "[3/6] Enabling IP forwarding..."
sysctl -w net.ipv4.ip_forward=1
if ! grep -q '^net.ipv4.ip_forward=1' /etc/sysctl.conf 2>/dev/null; then
    echo 'net.ipv4.ip_forward=1' >> /etc/sysctl.conf
fi

echo "[4/6] Starting WireGuard..."
systemctl enable wg-quick@${TUNNEL_INTERFACE}
systemctl start wg-quick@${TUNNEL_INTERFACE}

echo "[5/6] Configuring firewall..."
if command -v ufw &>/dev/null; then
    ufw allow ${TUNNEL_PORT}/udp
    ufw allow OpenSSH
    echo 'y' | ufw enable 2>/dev/null || true
elif command -v firewall-cmd &>/dev/null; then
    firewall-cmd --permanent --add-port=${TUNNEL_PORT}/udp
    firewall-cmd --reload
fi

echo "[6/6] Done."
echo ""
echo "Relay is running."
echo "  Tunnel IP: ${TUNNEL_RELAY_IP}"
echo "  Port: ${TUNNEL_PORT}"
echo "  Public Key: ${RELAY_PUBLIC_KEY}"
SETUP_SCRIPT
}

# ---------------------------------------------------------------------------
# --generate-script: output and exit
# ---------------------------------------------------------------------------
if [[ "$GENERATE_SCRIPT" == "true" ]]; then
    _build_setup_script
    echo "" >&2
    log_success "Setup script written to stdout. Pipe to a file or SSH."
    log_info "Example: tunnel-setup-relay.sh --generate-script | ssh root@YOUR_SERVER bash"
    exit 0
fi

# ---------------------------------------------------------------------------
# Provider: local (configure WireGuard on this machine)
# ---------------------------------------------------------------------------
if [[ "$PROVIDER" == "local" ]]; then
    log_info "Configuring this machine as the tunnel relay"
    require_cmd wg wg-quick

    # Run the setup script locally
    _build_setup_script | bash

    # Determine public IP
    RELAY_PUBLIC_IP="${RELAY_HOST:-}"
    if [[ -z "$RELAY_PUBLIC_IP" ]]; then
        RELAY_PUBLIC_IP="$(curl -4 -sS --max-time 5 https://ifconfig.me 2>/dev/null || \
                          curl -4 -sS --max-time 5 https://api.ipify.org 2>/dev/null || \
                          echo "UNKNOWN")"
    fi
    RELAY_ENDPOINT="${RELAY_PUBLIC_IP}:${TUNNEL_PORT}"

    registry_set_relay "$RELAY_ENDPOINT" "$RELAY_PUBLIC_KEY" "${TUNNEL_RELAY_IP}"

    audit_log "setup_relay" "success" "Configured local relay at $RELAY_ENDPOINT" \
        "$(jq -cn --arg ip "$RELAY_PUBLIC_IP" --arg mode "local" '{"relay_ip":$ip,"provider":$mode}')"

    echo ""
    echo "========================================"
    echo "  Tunnel Relay Configured (local)"
    echo "  Public IP: $RELAY_PUBLIC_IP"
    echo "  Relay Tunnel IP: ${TUNNEL_RELAY_IP}"
    echo "  Endpoint: $RELAY_ENDPOINT"
    echo "========================================"
    echo ""
    echo "Next steps:"
    echo "  1. On the target machine, run:"
    echo "     tunnel-join.sh $RELAY_ENDPOINT $RELAY_PUBLIC_KEY"
    echo ""
    echo "  2. Add peers with:"
    echo "     tunnel-add-peer.sh <peer-name>"
    echo ""
    log_success "Local relay setup complete"
    exit 0
fi

# ---------------------------------------------------------------------------
# Provider: ssh (configure an existing remote server via SSH)
# ---------------------------------------------------------------------------
if [[ "$PROVIDER" == "ssh" ]]; then
    if [[ -z "$RELAY_HOST" ]]; then
        log_error "RELAY_HOST required for --provider=ssh. Use --host=IP or export RELAY_HOST=IP"
        exit 1
    fi

    log_info "Configuring relay on $RELAY_HOST via SSH (user: $RELAY_SSH_USER)"
    require_cmd ssh

    # Run the setup script on the remote server
    _build_setup_script | ssh -o StrictHostKeyChecking=accept-new \
        "${RELAY_SSH_USER}@${RELAY_HOST}" bash

    RELAY_ENDPOINT="${RELAY_HOST}:${TUNNEL_PORT}"

    registry_set_relay "$RELAY_ENDPOINT" "$RELAY_PUBLIC_KEY" "${TUNNEL_RELAY_IP}"

    # Save server metadata
    jq -cn \
        --arg ip "$RELAY_HOST" \
        --arg user "$RELAY_SSH_USER" \
        --arg mode "ssh" \
        '{"relay_ip":$ip,"ssh_user":$user,"provider":$mode}' \
        > "${config_dir}/relay-server.json"

    audit_log "setup_relay" "success" "Configured relay on $RELAY_HOST via SSH" \
        "$(jq -cn --arg ip "$RELAY_HOST" --arg mode "ssh" '{"relay_ip":$ip,"provider":$mode}')"

    echo ""
    echo "========================================"
    echo "  Tunnel Relay Configured (SSH)"
    echo "  Server: $RELAY_HOST"
    echo "  Relay Tunnel IP: ${TUNNEL_RELAY_IP}"
    echo "  Endpoint: $RELAY_ENDPOINT"
    echo "========================================"
    echo ""
    echo "Next steps:"
    echo "  1. On the target machine, run:"
    echo "     tunnel-join.sh $RELAY_ENDPOINT $RELAY_PUBLIC_KEY"
    echo ""
    echo "  2. Add peers with:"
    echo "     tunnel-add-peer.sh <peer-name>"
    echo ""
    log_success "SSH relay setup complete"
    exit 0
fi

# ---------------------------------------------------------------------------
# Provider: digitalocean (provision a new Droplet)
# ---------------------------------------------------------------------------
if [[ "$PROVIDER" == "digitalocean" ]]; then
    require_env DO_API_TOKEN
    require_cmd wg

    log_info "Setting up tunnel relay Droplet"
    log_info "Region: ${TUNNEL_DO_REGION}, Size: ${TUNNEL_DO_SIZE}"

    # Check for doctl or fall back to curl
    USE_DOCTL=false
    if command -v doctl &>/dev/null; then
        USE_DOCTL=true
        log_info "Using doctl CLI"
    else
        log_info "doctl not found, using curl fallback"
        require_cmd curl
    fi

    # Upload SSH key (if not already present)
    SSH_KEY_PUB="${TUNNEL_SSH_KEY}.pub"
    if [[ ! -f "$SSH_KEY_PUB" ]]; then
        log_error "SSH public key not found: $SSH_KEY_PUB"
        log_error "Generate one: ssh-keygen -t ed25519 -f ${TUNNEL_SSH_KEY}"
        exit 1
    fi

    SSH_KEY_CONTENT="$(cat "$SSH_KEY_PUB")"
    SSH_KEY_FINGERPRINT="$(ssh-keygen -lf "$SSH_KEY_PUB" -E md5 | awk '{print $2}' | sed 's/MD5://')"

    log_info "SSH key fingerprint: $SSH_KEY_FINGERPRINT"

    DROPLET_TAG="${TUNNEL_APP_NAME}-tunnel"
    SSH_KEY_NAME="${TUNNEL_APP_NAME}-tunnel"

    if [[ "$USE_DOCTL" == "true" ]]; then
        if ! doctl compute ssh-key get "$SSH_KEY_FINGERPRINT" &>/dev/null; then
            doctl compute ssh-key create "$SSH_KEY_NAME" --public-key "$SSH_KEY_CONTENT"
            log_info "SSH key uploaded to DigitalOcean"
        fi
    else
        existing_keys="$(curl -sS -H "Authorization: Bearer ${DO_API_TOKEN}" \
            "https://api.digitalocean.com/v2/account/keys")"
        if ! echo "$existing_keys" | jq -e --arg fp "$SSH_KEY_FINGERPRINT" \
            '.ssh_keys[] | select(.fingerprint == $fp)' &>/dev/null; then
            curl -sS -X POST -H "Authorization: Bearer ${DO_API_TOKEN}" \
                -H "Content-Type: application/json" \
                -d "$(jq -cn --arg name "$SSH_KEY_NAME" --arg key "$SSH_KEY_CONTENT" \
                    '{"name":$name,"public_key":$key}')" \
                "https://api.digitalocean.com/v2/account/keys" >/dev/null
            log_info "SSH key uploaded to DigitalOcean"
        fi
    fi

    # Create cloud-init user data
    CLOUD_INIT="$(_build_setup_script)"

    # Create Droplet
    DROPLET_NAME="${TUNNEL_APP_NAME}-tunnel-relay"
    log_info "Creating Droplet: $DROPLET_NAME"

    if [[ "$USE_DOCTL" == "true" ]]; then
        DROPLET_ID="$(doctl compute droplet create "$DROPLET_NAME" \
            --region "$TUNNEL_DO_REGION" \
            --size "$TUNNEL_DO_SIZE" \
            --image "ubuntu-24-04-x64" \
            --ssh-keys "$SSH_KEY_FINGERPRINT" \
            --user-data "$CLOUD_INIT" \
            --tag-name "$DROPLET_TAG" \
            --wait \
            --format ID \
            --no-header)"

        DROPLET_IP="$(doctl compute droplet get "$DROPLET_ID" --format PublicIPv4 --no-header)"
    else
        CREATE_RESPONSE="$(curl -sS -X POST -H "Authorization: Bearer ${DO_API_TOKEN}" \
            -H "Content-Type: application/json" \
            -d "$(jq -cn \
                --arg name "$DROPLET_NAME" \
                --arg region "$TUNNEL_DO_REGION" \
                --arg size "$TUNNEL_DO_SIZE" \
                --arg image "ubuntu-24-04-x64" \
                --arg fp "$SSH_KEY_FINGERPRINT" \
                --arg ud "$CLOUD_INIT" \
                --arg tag "$DROPLET_TAG" \
                '{
                    "name": $name,
                    "region": $region,
                    "size": $size,
                    "image": $image,
                    "ssh_keys": [$fp],
                    "user_data": $ud,
                    "tags": [$tag]
                }')" \
            "https://api.digitalocean.com/v2/droplets")"

        DROPLET_ID="$(echo "$CREATE_RESPONSE" | jq -r '.droplet.id')"
        if [[ "$DROPLET_ID" == "null" || -z "$DROPLET_ID" ]]; then
            log_error "Failed to create Droplet"
            log_error "Response: $CREATE_RESPONSE"
            exit 1
        fi

        log_info "Droplet $DROPLET_ID created, waiting for IP..."

        # Poll for IP (up to 120s)
        DROPLET_IP=""
        for _ in $(seq 1 24); do
            sleep 5
            DROPLET_INFO="$(curl -sS -H "Authorization: Bearer ${DO_API_TOKEN}" \
                "https://api.digitalocean.com/v2/droplets/${DROPLET_ID}")"
            DROPLET_IP="$(echo "$DROPLET_INFO" | jq -r '.droplet.networks.v4[] | select(.type == "public") | .ip_address' | head -1)"
            if [[ -n "$DROPLET_IP" && "$DROPLET_IP" != "null" ]]; then
                break
            fi
        done

        if [[ -z "$DROPLET_IP" || "$DROPLET_IP" == "null" ]]; then
            log_error "Timed out waiting for Droplet IP"
            exit 1
        fi
    fi

    RELAY_ENDPOINT="${DROPLET_IP}:${TUNNEL_PORT}"
    log_success "Droplet ready: $DROPLET_IP"

    # Save relay info
    registry_set_relay "$RELAY_ENDPOINT" "$RELAY_PUBLIC_KEY" "${TUNNEL_RELAY_IP}"

    # Save Droplet metadata
    jq -cn \
        --arg id "$DROPLET_ID" \
        --arg ip "$DROPLET_IP" \
        --arg name "$DROPLET_NAME" \
        --arg region "$TUNNEL_DO_REGION" \
        --arg mode "digitalocean" \
        '{"droplet_id":$id,"ip":$ip,"name":$name,"region":$region,"provider":$mode}' \
        > "${config_dir}/relay-server.json"

    audit_log "setup_relay" "success" "Provisioned relay Droplet $DROPLET_IP" \
        "$(jq -cn --arg ip "$DROPLET_IP" --arg id "$DROPLET_ID" --arg mode "digitalocean" \
            '{"droplet_ip":$ip,"droplet_id":$id,"provider":$mode}')"

    echo ""
    echo "========================================"
    echo "  Tunnel Relay Provisioned"
    echo "  Droplet IP: $DROPLET_IP"
    echo "  Relay Tunnel IP: ${TUNNEL_RELAY_IP}"
    echo "  Endpoint: $RELAY_ENDPOINT"
    echo "========================================"
    echo ""
    echo "Next steps:"
    echo "  1. On the target machine, run:"
    echo "     tunnel-join.sh $RELAY_ENDPOINT $RELAY_PUBLIC_KEY"
    echo ""
    echo "  2. Add peers with:"
    echo "     tunnel-add-peer.sh <peer-name>"
    echo ""
    log_success "Relay setup complete"
    exit 0
fi

log_error "Unknown provider: $PROVIDER"
log_error "Valid providers: digitalocean, ssh, local"
exit 1