#!/bin/bash
# Dedicated private-LAN WireGuard hub. Does not enable an Internet exit/NAT.
# Usage: bash .../17-install-wireguard.sh <ssh-host> <public-ip-or-host> <private-output-dir>
set -euo pipefail
umask 077
unset HTTPS_PROXY HTTP_PROXY https_proxy http_proxy ALL_PROXY all_proxy 2>/dev/null || true
host="${1:?SSH host required}"
endpoint="${2:?Public endpoint required}"
output_dir="${3:?Private output directory required}"
[[ "$host" =~ ^[a-z0-9-]+$ && "$endpoint" =~ ^[a-zA-Z0-9.-]+$ ]] || exit 2
mkdir -p "$output_dir"
chmod 700 "$output_dir"

ssh -o BatchMode=yes -o ConnectTimeout=15 "$host" 'set -e
    if ! command -v wg >/dev/null || ! command -v wg-quick >/dev/null; then
        sudo -n env DEBIAN_FRONTEND=noninteractive NEEDRESTART_MODE=l apt-get update -qq
        sudo -n env DEBIAN_FRONTEND=noninteractive NEEDRESTART_MODE=l apt-get install -y --no-install-recommends wireguard-tools
    fi'

ssh -o BatchMode=yes -o ConnectTimeout=15 "$host" sudo -n python3 - "$endpoint" <<'PY'
import datetime
import ipaddress
import json
import os
from pathlib import Path
import subprocess
import sys

os.umask(0o077)
endpoint = sys.argv[1]
interface = 'wg-njcq'
subnet = '10.203.0.0/24'
port = 51820
root = Path('/etc/wireguard')
config = root / (interface + '.conf')
clients = root / 'njcq-clients.json'
if config.exists():
    if not clients.exists():
        raise SystemExit('Existing interface has no managed client inventory; refusing replacement')
    print('Existing managed WireGuard hub preserved')
    raise SystemExit(0)
routes = json.loads(subprocess.check_output(['ip', '-j', 'route']))
for route in routes:
    dst = route.get('dst', 'default')
    if dst != 'default' and ipaddress.ip_network(dst, strict=False).overlaps(ipaddress.ip_network(subnet)):
        raise SystemExit('Subnet overlaps an existing route: ' + dst)
stamp = datetime.datetime.now().strftime('%Y%m%d-%H%M%S')
backup = Path('/root/backups') / ('before-wireguard-config-' + stamp)
backup.mkdir(parents=True, mode=0o700)
(backup / 'iptables.v4').write_bytes(subprocess.check_output(['iptables-save']))
(backup / 'forwarding.txt').write_bytes(subprocess.check_output(['sysctl', 'net.ipv4.ip_forward']))
root.mkdir(mode=0o700, exist_ok=True)
root.chmod(0o700)

def keypair():
    private = subprocess.check_output(['wg', 'genkey'], text=True).strip()
    public = subprocess.check_output(['wg', 'pubkey'], input=private, text=True).strip()
    return private, public

server_private, server_public = keypair()
devices = []
for name, address, incoming in [('iPad Pro', '10.203.0.2', False), ('Mac Studio', '10.203.0.3', True)]:
    private, public = keypair()
    psk = subprocess.check_output(['wg', 'genpsk'], text=True).strip()
    devices.append({'name': name, 'publicKey': public, 'profile': {
        'server': endpoint, 'port': port, 'serverPublicKey': server_public,
        'privateKey': private, 'presharedKey': psk, 'address': address,
        'subnet': subnet, 'mtu': 1280, 'keepalive': 25, 'acceptIncoming': incoming,
    }})

# Private forwarding is allowed only between peers of this interface.
# The other interfaces, Docker and Tailscale keep their existing policies.
up = [
    f'iptables -I INPUT 1 -p udp --dport {port} -m comment --comment proxy-station-wireguard -j ACCEPT',
    f'iptables -I FORWARD 1 -i %i -o %i -s {subnet} -d {subnet} -j ACCEPT',
    'iptables -I FORWARD 2 -i %i -j DROP',
    'iptables -I FORWARD 3 -o %i -j DROP',
]
down = [
    f'iptables -D INPUT -p udp --dport {port} -m comment --comment proxy-station-wireguard -j ACCEPT',
    f'iptables -D FORWARD -i %i -o %i -s {subnet} -d {subnet} -j ACCEPT',
    'iptables -D FORWARD -i %i -j DROP',
    'iptables -D FORWARD -o %i -j DROP',
]
lines = ['[Interface]', 'Address = 10.203.0.1/24', f'ListenPort = {port}',
         f'PrivateKey = {server_private}', 'MTU = 1280']
lines += ['PostUp = ' + rule for rule in up]
lines += ['PostDown = ' + rule for rule in down]
for device in devices:
    profile = device['profile']
    lines += ['', '# ' + device['name'], '[Peer]', 'PublicKey = ' + device['publicKey'],
              'PresharedKey = ' + profile['presharedKey'], 'AllowedIPs = ' + profile['address'] + '/32']
config.write_text('\n'.join(lines) + '\n')
config.chmod(0o600)
clients.write_text(json.dumps({'interface': interface, 'serverPublicKey': server_public, 'devices': devices}, indent=2))
clients.chmod(0o600)
Path('/etc/sysctl.d/90-proxy-station-wireguard.conf').write_text('net.ipv4.ip_forward = 1\n')
subprocess.run(['sysctl', '-w', 'net.ipv4.ip_forward=1'], check=True)
try:
    subprocess.run(['systemctl', 'enable', '--now', 'wg-quick@' + interface], check=True)
except subprocess.CalledProcessError:
    subprocess.run(['systemctl', 'disable', '--now', 'wg-quick@' + interface], check=False)
    raise
print('WireGuard hub installed:', interface, subnet, 'UDP', port)
print('Backup:', backup)
print('Devices:', ', '.join(d['name'] for d in devices))
PY

ssh -o BatchMode=yes -o ConnectTimeout=15 "$host" 'sudo -n cat /etc/wireguard/njcq-clients.json' > "$output_dir/clients.json"
test -s "$output_dir/clients.json"
chmod 600 "$output_dir/clients.json"
ssh -o BatchMode=yes -o ConnectTimeout=15 "$host" 'sudo -n wg show wg-njcq; systemctl is-active wg-quick@wg-njcq tailscaled'