#!/usr/bin/env bash
set -euo pipefail

MODE="${1:-}"
STATE_DIR="/tmp/soria-wireguard-lab"

SITE_NAMESPACES=("wg-a-host" "wg-a-rtr" "wg-b-host" "wg-b-rtr")
REMOTE_NAMESPACES=("wg-hub" "wg-admin" "wg-road")

usage() {
  cat <<'EOF'
Usage:
  sudo ./module4b-wireguard-lab.sh site
  sudo ./module4b-wireguard-lab.sh remote
  sudo ./module4b-wireguard-lab.sh rotate-road
  sudo ./module4b-wireguard-lab.sh status
  sudo ./module4b-wireguard-lab.sh cleanup

site        : two routed LANs connected by a WireGuard site-to-site tunnel.
remote      : one road-warrior client reaching an administration LAN.
rotate-road : replace the road-warrior key pair and update the hub peer.
status      : show namespaces, routes and WireGuard state.
cleanup     : stop lab processes, delete namespaces and remove lab keys.
EOF
}

require_root() {
  if [[ ${EUID} -ne 0 ]]; then
    echo "Exécute ce script avec sudo." >&2
    exit 1
  fi
}

require_commands() {
  local missing=0
  for command in ip wg sysctl; do
    if ! command -v "$command" >/dev/null 2>&1; then
      echo "Commande manquante : $command" >&2
      missing=1
    fi
  done
  if [[ $missing -ne 0 ]]; then
    exit 1
  fi
}

kill_namespace_processes() {
  local ns="$1"
  local pids
  pids="$(ip netns pids "$ns" 2>/dev/null || true)"
  if [[ -n "$pids" ]]; then
    kill $pids 2>/dev/null || true
    sleep 0.2
  fi
}

cleanup_namespaces() {
  local ns
  for ns in "${SITE_NAMESPACES[@]}" "${REMOTE_NAMESPACES[@]}"; do
    kill_namespace_processes "$ns"
    ip netns del "$ns" 2>/dev/null || true
  done
}

cleanup() {
  cleanup_namespaces
  rm -rf "$STATE_DIR"
}

create_ns() {
  local ns="$1"
  ip netns add "$ns"
  ip -n "$ns" link set lo up
}

link_namespaces() {
  local ns_a="$1" root_a="$2" if_a="$3"
  local ns_b="$4" root_b="$5" if_b="$6"

  ip link add "$root_a" type veth peer name "$root_b"
  ip link set "$root_a" netns "$ns_a"
  ip link set "$root_b" netns "$ns_b"
  ip -n "$ns_a" link set "$root_a" name "$if_a"
  ip -n "$ns_b" link set "$root_b" name "$if_b"
  ip -n "$ns_a" link set "$if_a" up
  ip -n "$ns_b" link set "$if_b" up
}

generate_keypair() {
  local name="$1"
  umask 077
  wg genkey | tee "${STATE_DIR}/${name}.key" | wg pubkey > "${STATE_DIR}/${name}.pub"
}

pub() {
  cat "${STATE_DIR}/$1.pub"
}

setup_site() {
  cleanup
  mkdir -p "$STATE_DIR"
  chmod 700 "$STATE_DIR"

  for ns in "${SITE_NAMESPACES[@]}"; do
    create_ns "$ns"
  done

  link_namespaces "wg-a-host" ahost eth0 "wg-a-rtr" arlan lan0
  link_namespaces "wg-b-host" bhost eth0 "wg-b-rtr" brlan lan0
  link_namespaces "wg-a-rtr" awan wan0 "wg-b-rtr" bwan wan0

  ip -n wg-a-host addr add 10.60.10.10/24 dev eth0
  ip -n wg-a-rtr  addr add 10.60.10.1/24 dev lan0
  ip -n wg-b-host addr add 10.60.20.10/24 dev eth0
  ip -n wg-b-rtr  addr add 10.60.20.1/24 dev lan0

  ip -n wg-a-rtr addr add 192.0.2.1/30 dev wan0
  ip -n wg-b-rtr addr add 192.0.2.2/30 dev wan0

  ip -n wg-a-host route add default via 10.60.10.1
  ip -n wg-b-host route add default via 10.60.20.1

  ip netns exec wg-a-rtr sysctl -qw net.ipv4.ip_forward=1
  ip netns exec wg-b-rtr sysctl -qw net.ipv4.ip_forward=1

  generate_keypair site-a
  generate_keypair site-b

  ip -n wg-a-rtr link add wg0 type wireguard
  ip -n wg-b-rtr link add wg0 type wireguard

  ip netns exec wg-a-rtr wg set wg0 \
    private-key "${STATE_DIR}/site-a.key" \
    listen-port 51820 \
    peer "$(pub site-b)" \
    endpoint 192.0.2.2:51820 \
    allowed-ips 10.255.0.2/32,10.60.20.0/24

  ip netns exec wg-b-rtr wg set wg0 \
    private-key "${STATE_DIR}/site-b.key" \
    listen-port 51820 \
    peer "$(pub site-a)" \
    endpoint 192.0.2.1:51820 \
    allowed-ips 10.255.0.1/32,10.60.10.0/24

  ip -n wg-a-rtr addr add 10.255.0.1/30 dev wg0
  ip -n wg-b-rtr addr add 10.255.0.2/30 dev wg0
  ip -n wg-a-rtr link set wg0 up
  ip -n wg-b-rtr link set wg0 up

  ip -n wg-a-rtr route add 10.60.20.0/24 dev wg0
  ip -n wg-b-rtr route add 10.60.10.0/24 dev wg0

  cat <<'EOF'
WireGuard site-to-site prêt.

Tests :
  sudo ip netns exec wg-a-host ping -c 3 10.60.20.10
  sudo ip netns exec wg-a-rtr wg show
  sudo ip netns exec wg-a-rtr tcpdump -ni wan0 udp port 51820
EOF
}

setup_remote() {
  cleanup
  mkdir -p "$STATE_DIR"
  chmod 700 "$STATE_DIR"

  for ns in "${REMOTE_NAMESPACES[@]}"; do
    create_ns "$ns"
  done

  link_namespaces "wg-road" roadwan eth0 "wg-hub" hubwan wan0
  link_namespaces "wg-admin" admineth eth0 "wg-hub" hublan lan0

  ip -n wg-road  addr add 198.51.100.2/30 dev eth0
  ip -n wg-hub   addr add 198.51.100.1/30 dev wan0
  ip -n wg-admin addr add 10.70.10.10/24 dev eth0
  ip -n wg-hub   addr add 10.70.10.1/24 dev lan0

  ip -n wg-admin route add 10.254.0.0/24 via 10.70.10.1
  ip netns exec wg-hub sysctl -qw net.ipv4.ip_forward=1

  generate_keypair hub
  generate_keypair road

  ip -n wg-hub  link add wg0 type wireguard
  ip -n wg-road link add wg0 type wireguard

  ip netns exec wg-hub wg set wg0 \
    private-key "${STATE_DIR}/hub.key" \
    listen-port 51820 \
    peer "$(pub road)" \
    allowed-ips 10.254.0.10/32

  ip netns exec wg-road wg set wg0 \
    private-key "${STATE_DIR}/road.key" \
    peer "$(pub hub)" \
    endpoint 198.51.100.1:51820 \
    allowed-ips 10.254.0.1/32,10.70.10.0/24 \
    persistent-keepalive 25

  ip -n wg-hub  addr add 10.254.0.1/24 dev wg0
  ip -n wg-road addr add 10.254.0.10/24 dev wg0
  ip -n wg-hub  link set wg0 up
  ip -n wg-road link set wg0 up

  ip -n wg-road route add 10.70.10.0/24 dev wg0

  cat <<'EOF'
WireGuard accès distant prêt.

Tests :
  sudo ip netns exec wg-road ping -c 3 10.70.10.10
  sudo ip netns exec wg-hub wg show
  sudo ls -l /tmp/soria-wireguard-lab
EOF
}

rotate_road() {
  if ! ip netns list | grep -q '^wg-hub'; then
    echo "Lance d'abord le mode remote." >&2
    exit 1
  fi
  if [[ ! -f "${STATE_DIR}/road.pub" ]]; then
    echo "Clé road introuvable." >&2
    exit 1
  fi

  local old_pub new_pub
  old_pub="$(pub road)"

  umask 077
  wg genkey | tee "${STATE_DIR}/road.next.key" | wg pubkey > "${STATE_DIR}/road.next.pub"
  new_pub="$(cat "${STATE_DIR}/road.next.pub")"

  ip netns exec wg-hub wg set wg0 peer "$old_pub" remove
  ip netns exec wg-road wg set wg0 private-key "${STATE_DIR}/road.next.key"
  ip netns exec wg-hub wg set wg0 peer "$new_pub" allowed-ips 10.254.0.10/32

  mv "${STATE_DIR}/road.next.key" "${STATE_DIR}/road.key"
  mv "${STATE_DIR}/road.next.pub" "${STATE_DIR}/road.pub"

  echo "Clé road remplacée. Vérifie le nouveau handshake."
}

status() {
  local ns
  for ns in "${SITE_NAMESPACES[@]}" "${REMOTE_NAMESPACES[@]}"; do
    if ip netns list | grep -q "^${ns}\b"; then
      echo
      echo "[$ns]"
      ip -n "$ns" -brief address
      ip -n "$ns" route
      ip netns exec "$ns" wg show 2>/dev/null || true
    fi
  done
}

require_root

case "$MODE" in
  site)
    require_commands
    setup_site
    ;;
  remote)
    require_commands
    setup_remote
    ;;
  rotate-road)
    require_commands
    rotate_road
    ;;
  status)
    require_commands
    status
    ;;
  cleanup)
    if ! command -v ip >/dev/null 2>&1; then
      echo "Commande manquante : ip" >&2
      exit 1
    fi
    cleanup
    echo "Laboratoire supprimé."
    ;;
  *)
    usage
    exit 1
    ;;
esac