#!/bin/bash
################################################################################
# Copyright (c) 2026 ModalAI, Inc. All rights reserved.
#
# voxl-wg-keygen - Generate keys and write a complete /etc/wireguard/wg0.conf
#                  in one shot. Prints a ready-to-paste server-side command
#                  for registering this device.
################################################################################

set -e

ADDRESS=""
FORCE=0
NAME=""
SERVER_CONF="/etc/voxl-wireguard/server.conf"
PRIV=/etc/wireguard/privatekey
PUB=/etc/wireguard/publickey
WG_CONF=/etc/wireguard/wg0.conf

print_usage() {
	cat <<EOF
voxl-wg-keygen - Set up WireGuard on this device

Usage:
  voxl-wg-keygen --address <CIDR> [--name <NAME>] [--force]

Options:
  --address <CIDR>   This device's VPN IP, e.g. 10.8.0.2/24
                     Convention:
                       10.8.0.2/24   for ground stations
                       10.8.0.11/24+ for drones
  --name <NAME>      Optional label for this peer. Used as a comment
                     above the [Peer] block on the server side. Defaults
                     to this device's hostname.
  --force            Overwrite existing keys and config.
                     Requires re-registering this device with the server.
  -h, --help         Show this help.

Server connection settings are read from $SERVER_CONF.

Examples:
  voxl-wg-keygen --address 10.8.0.2/24                    # ground station
  voxl-wg-keygen --address 10.8.0.11/24 --name drone-1    # drone
EOF
}

while [ $# -gt 0 ]; do
	case "$1" in
		--address) ADDRESS="$2"; shift 2 ;;
		--name)    NAME="$2";    shift 2 ;;
		--force)   FORCE=1;      shift ;;
		-h|--help) print_usage;  exit 0 ;;
		*)
			echo "ERROR: unknown argument: $1" >&2
			print_usage >&2
			exit 1
			;;
	esac
done

[ "$(id -u)" -eq 0 ] || { echo "ERROR: must run as root" >&2; exit 1; }

if [ -z "$ADDRESS" ]; then
	echo "ERROR: --address is required" >&2
	print_usage >&2
	exit 1
fi

# Validate address format
if ! echo "$ADDRESS" | grep -qE '^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+/[0-9]+$'; then
	echo "ERROR: --address must be CIDR like 10.8.0.2/24 (got: $ADDRESS)" >&2
	exit 1
fi

# Default name to the device hostname
if [ -z "$NAME" ]; then
	NAME=$(hostname 2>/dev/null || echo "unnamed")
fi

# Load server config
if [ ! -r "$SERVER_CONF" ]; then
	echo "ERROR: $SERVER_CONF not readable" >&2
	exit 1
fi
# shellcheck disable=SC1090
. "$SERVER_CONF"

[ -n "$SERVER_PUBKEY" ]   || { echo "ERROR: SERVER_PUBKEY not set in $SERVER_CONF" >&2; exit 1; }
[ -n "$SERVER_ENDPOINT" ] || { echo "ERROR: SERVER_ENDPOINT not set in $SERVER_CONF" >&2; exit 1; }
SERVER_SUBNET="${SERVER_SUBNET:-10.8.0.0/24}"
MTU="${MTU:-1380}"
PERSISTENT_KEEPALIVE="${PERSISTENT_KEEPALIVE:-25}"

mkdir -p /etc/wireguard
chmod 700 /etc/wireguard

# Generate or reuse keypair
if [ -e "$PRIV" ] && [ "$FORCE" -eq 0 ]; then
	echo "Existing keypair found at $PRIV — reusing"
	echo "(use --force to regenerate; this requires re-registering with the server)"
else
	umask 077
	wg genkey | tee "$PRIV" | wg pubkey > "$PUB"
	chmod 600 "$PRIV" "$PUB"
	echo "New keypair written:"
	echo "  $PRIV"
	echo "  $PUB"
fi

PRIVATE_KEY=$(cat "$PRIV")
PUBLIC_KEY=$(cat "$PUB")

# Refuse to clobber an existing wg0.conf unless --force
if [ -e "$WG_CONF" ] && [ "$FORCE" -eq 0 ]; then
	echo "ERROR: $WG_CONF already exists. Use --force to overwrite." >&2
	exit 1
fi

# Write wg0.conf
umask 077
cat > "$WG_CONF" <<EOF
# Generated by voxl-wg-keygen on $(date -u +%Y-%m-%dT%H:%M:%SZ)
# Re-run with --force to regenerate.

[Interface]
Address = $ADDRESS
PrivateKey = $PRIVATE_KEY
MTU = $MTU

[Peer]
PublicKey = $SERVER_PUBKEY
Endpoint = $SERVER_ENDPOINT
AllowedIPs = $SERVER_SUBNET
PersistentKeepalive = $PERSISTENT_KEEPALIVE
EOF
chmod 600 "$WG_CONF"

# Strip CIDR from address for the server-side command
ADDRESS_NO_CIDR="${ADDRESS%/*}"

echo ""
echo "$WG_CONF written"
echo "  Address:   $ADDRESS"
echo "  Server:    $SERVER_ENDPOINT"
echo "  Subnet:    $SERVER_SUBNET"
echo "  Name:      $NAME"
echo ""
echo "==============================================================="
echo "On the WireGuard server, run:"
echo ""
echo "    sudo voxl-wg-server-add-peer \\"
echo "        --pubkey $PUBLIC_KEY \\"
echo "        --address $ADDRESS_NO_CIDR \\"
echo "        --name $NAME"
echo ""
echo "Then on this device:"
echo ""
echo "    systemctl enable --now voxl-wireguard"
echo "==============================================================="
