#!/bin/bash
# Manage the WireGuard VPN network namespace.
# Usage: vpn-ns up | down | status
set -euo pipefail

NS="ns-gemma"
WG_CONF="/etc/wireguard/gemma.conf"
LOCKFILE="/run/vpn-ns-gemma.lock"

up() {
  # Use a lock to prevent race conditions when multiple processes
  # try to create the namespace concurrently
  exec 200>"$LOCKFILE"
  flock 200

  if ip netns list | grep -q "^${NS}"; then
    echo "Namespace ${NS} already exists"
    return 0
  fi

  if [ ! -f "$WG_CONF" ]; then
    echo "ERROR: WireGuard config not found at ${WG_CONF}" >&2
    exit 1
  fi

  echo "Creating network namespace ${NS}..."

  # Create namespace and bring up loopback
  ip netns add "$NS"
  ip netns exec "$NS" ip link set lo up

  # Create WireGuard interface inside namespace
  ip link add wg-gemma type wireguard
  ip link set wg-gemma netns "$NS"

  # Extract Address from config
  local addr
  addr=$(grep -i '^\s*Address\s*=' "$WG_CONF" | head -1 | sed 's/.*=\s*//' | tr -d ' ')

  # Strip config down to what `wg setconf` expects (no Address, DNS, etc.)
  # Also resolve any Endpoint hostnames to IPs using host DNS (the namespace
  # has no network yet, so DNS won't work inside it at this point).
  local tmpconf
  tmpconf=$(mktemp)
  grep -iv '^\s*\(Address\|DNS\|MTU\|Table\|PreUp\|PostUp\|PreDown\|PostDown\|SaveConfig\)\s*=' "$WG_CONF" > "$tmpconf"

  # Resolve Endpoint hostnames to IPs (e.g. "vpn.example.com:51820" -> "1.2.3.4:51820")
  local endpoint_line endpoint_host endpoint_port resolved_ip
  endpoint_line=$(grep -i '^\s*Endpoint\s*=' "$tmpconf" | head -1 | sed 's/.*=\s*//' | tr -d ' ' || true)
  if [ -n "$endpoint_line" ]; then
    endpoint_host="${endpoint_line%:*}"
    endpoint_port="${endpoint_line##*:}"
    # Only resolve if it's not already an IP
    if ! echo "$endpoint_host" | grep -qE '^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$'; then
      resolved_ip=$(getent hosts "$endpoint_host" | awk '{print $1; exit}')
      if [ -n "$resolved_ip" ]; then
        sed -i "s|${endpoint_host}:${endpoint_port}|${resolved_ip}:${endpoint_port}|g" "$tmpconf"
      else
        echo "ERROR: Could not resolve endpoint hostname: ${endpoint_host}" >&2
        rm -f "$tmpconf"
        ip netns del "$NS"
        exit 1
      fi
    fi
  fi

  ip netns exec "$NS" wg setconf wg-gemma "$tmpconf"
  rm -f "$tmpconf"

  # Apply MTU if specified in config
  local mtu
  mtu=$(grep -i '^\s*MTU\s*=' "$WG_CONF" | head -1 | sed 's/.*=\s*//' | tr -d ' ' || true)
  if [ -n "$mtu" ]; then
    ip netns exec "$NS" ip link set wg-gemma mtu "$mtu"
  fi

  # Assign address(es) and bring up interface
  for a in $(echo "$addr" | tr ',' ' '); do
    ip netns exec "$NS" ip addr add "$a" dev wg-gemma
  done
  ip netns exec "$NS" ip link set wg-gemma up

  # Add routes for AllowedIPs
  local allowed_ips
  allowed_ips=$(grep -i '^\s*AllowedIPs\s*=' "$WG_CONF" | sed 's/.*=\s*//' | tr ',' '\n' | tr -d ' ')
  for cidr in $allowed_ips; do
    # For default routes (0.0.0.0/0 or ::/0), use "default" syntax
    if [ "$cidr" = "0.0.0.0/0" ]; then
      ip netns exec "$NS" ip route add default dev wg-gemma 2>/dev/null || true
    elif [ "$cidr" = "::/0" ]; then
      ip netns exec "$NS" ip -6 route add default dev wg-gemma 2>/dev/null || true
    else
      ip netns exec "$NS" ip route add "$cidr" dev wg-gemma 2>/dev/null || true
    fi
  done

  # Set up DNS inside the namespace
  mkdir -p "/etc/netns/${NS}"
  local dns
  dns=$(grep -i '^\s*DNS\s*=' "$WG_CONF" | head -1 | sed 's/.*=\s*//' | tr -d ' ' || true)
  if [ -n "$dns" ]; then
    for d in $(echo "$dns" | tr ',' ' '); do
      echo "nameserver $d"
    done > "/etc/netns/${NS}/resolv.conf"
  else
    cp /etc/resolv.conf "/etc/netns/${NS}/resolv.conf"
  fi

  echo "VPN namespace ${NS} is up"
}

down() {
  if ! ip netns list | grep -q "^${NS}"; then
    echo "Namespace ${NS} does not exist"
    return 0
  fi

  echo "Removing network namespace ${NS}..."
  ip netns del "$NS"
  rm -rf "/etc/netns/${NS}"
  echo "VPN namespace ${NS} is down"
}

status() {
  if ! ip netns list | grep -q "^${NS}"; then
    echo "VPN namespace: down"
    return 1
  fi

  echo "VPN namespace: up"
  ip netns exec "$NS" wg show 2>/dev/null || echo "(no WireGuard info)"
  echo ""
  echo "Routes inside namespace:"
  ip netns exec "$NS" ip route 2>/dev/null
}

case "${1:-}" in
  up) up ;;
  down) down ;;
  status) status ;;
  *)
    echo "Usage: vpn-ns up|down|status" >&2
    exit 1
    ;;
esac