#!/bin/bash
# Prove that VyOS hands each device the address UniFi reserved for it.
#
# The question this answers is narrow and important: 30 of the 31 UniFi
# reservations sit INSIDE the DHCP pool (LoT's pool is 10.0.0.11-10.0.1.254 and
# only 10.0.0.2 falls outside it). UniFi's dhcpd tolerates that. VyOS uses kea,
# and whether kea honours in-pool host reservations decides whether the cutover
# silently renumbers 30 devices. That is not something to predict.
#
# Method: boot throwaway VMs whose MAC is a REAL production MAC, on the sim
# VLAN, and check the address they are given. MACs are the one piece of
# production config that transplants verbatim -- the subnet is rewritten, the
# MAC is not -- which is what makes this a real test rather than a rehearsal.
#
# Safe: the ovs-labsim bridge contains only internal ports and VM taps, with no
# physical NIC, so a production MAC here cannot reach or confuse the real LAN.
# Verified with `ovs-vsctl show` before this script was written.
#
#   ./labsim-dhcp-test.sh              run the standard cases
#   ./labsim-dhcp-test.sh --keep       leave the VMs up for inspection
#   ./labsim-dhcp-test.sh --clean      just remove any leftover test VMs
set -uo pipefail

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
source "$SCRIPT_DIR/lib.sh"

ROUTER_IP="${ROUTER_IP:-172.31.1.1}"
ROUTER_PW="${ROUTER_PW:-vyos}"
TEST_VLAN="${TEST_VLAN:-10}"
BOOT_WAIT="${BOOT_WAIT:-150}"
TAG="labsim-dhcptest"

# mac|expected|why. "POOL" means: must get an address from the pool and must
# NOT get any reserved address -- the negative case that stops a pass from
# meaning merely "DHCP works".
CASES=(
  "f8:0d:ac:90:65:c6|172.31.10.46|printer1 - reservation inside the pool"
  "1c:69:20:7f:bc:77|172.31.11.67|sonoff-matter - in-pool AND across the /23 boundary"
  "34:e1:d1:80:29:ce|172.31.10.2|Hubitat - the one reservation OUTSIDE the pool"
  "52:54:00:ab:cd:ef|POOL|unreserved MAC - must get a pool address, not a reserved one"
)

vm_of() { echo "${TAG}-$(echo "$1" | tr -d ':')"; }

cleanup_vms() {
  local n=0
  while read -r vm; do
    [ -z "$vm" ] && continue
    virsh_q destroy "$vm" >/dev/null 2>&1
    virsh_q undefine "$vm" --remove-all-storage >/dev/null 2>&1
    n=$((n + 1))
  done < <(virsh_q list --all --name 2>/dev/null | grep "^${TAG}-" || true)
  [ "$n" -gt 0 ] && log "removed $n test VM(s)"
  sudo rm -f "$IMG_DIR/${TAG}-"*.qcow2 "$IMG_DIR/${TAG}-"*-seed.iso 2>/dev/null
  return 0
}

# A seed that asks for DHCP instead of taking a static address. Alpine's
# cloud-init ignores network-config here (verified previously and documented in
# README), so /etc/network/interfaces is what actually takes effect.
build_dhcp_seed() {
  local iso="$1" vm="$2" pubkey="$3"
  local tmp; tmp="$(mktemp -d)"
  cat > "$tmp/meta-data" <<EOF
instance-id: $vm
local-hostname: $vm
EOF
  cat > "$tmp/user-data" <<EOF
#cloud-config
hostname: $vm
users:
  - name: alpine
    shell: /bin/ash
    lock_passwd: false
    plain_text_passwd: labsim
    ssh_authorized_keys:
      - $pubkey
ssh_authorized_keys:
  - $pubkey
disable_root: false
chpasswd:
  list: |
    root:labsim
  expire: false
write_files:
  - path: /etc/network/interfaces
    content: |
      auto lo
      iface lo inet loopback
      auto eth0
      iface eth0 inet dhcp
runcmd:
  - [ sh, -c, "ifdown eth0 2>/dev/null; ifup eth0 || udhcpc -i eth0 -q || true" ]
EOF
  python3 - "$tmp/user-data" <<'PY' || die "generated user-data is not valid YAML"
import sys, yaml
yaml.safe_load(open(sys.argv[1]).read().split("#cloud-config",1)[1])
PY
  sudo genisoimage -quiet -output "$iso" -volid cidata -joliet -rock \
    "$tmp/user-data" "$tmp/meta-data" >/dev/null 2>&1 || die "seed build failed"
  rm -rf "$tmp"
}

router() {
  timeout 30 sshpass -p "$ROUTER_PW" ssh -o StrictHostKeyChecking=no \
    -o BatchMode=no -o ConnectTimeout=8 "vyos@$ROUTER_IP" "$@" 2>/dev/null
}

# --- argument handling ----------------------------------------------------
KEEP=0
case "${1:-}" in
  --clean) cleanup_vms; exit 0 ;;
  --keep)  KEEP=1 ;;
  "")      ;;
  *)       die "usage: $0 [--keep|--clean]" ;;
esac

command -v sshpass >/dev/null || die "sshpass required"
require_tools
[ -f "$BASE_IMAGE" ] || die "base image missing: $BASE_IMAGE (run labsim-up.sh first)"

log "checking the router is serving DHCP..."
subnets=$(router '/opt/vyatta/bin/vyatta-op-cmd-wrapper show configuration commands | grep -c subnet-id')
maps=$(router '/opt/vyatta/bin/vyatta-op-cmd-wrapper show configuration commands | grep -c "static-mapping .* mac"')
log "  router has ${subnets:-0} subnets and ${maps:-0} static-mappings"
[ "${maps:-0}" -gt 0 ] || die "router has no static-mappings -- apply the generated config first"

cleanup_vms
SSH_PUB="$(find_ssh_pubkey)"
sudo mkdir -p "$IMG_DIR"

# --- boot one VM per case -------------------------------------------------
for c in "${CASES[@]}"; do
  IFS='|' read -r mac expected why <<<"$c"
  vm="$(vm_of "$mac")"
  disk="$IMG_DIR/${vm}.qcow2"; seed="$IMG_DIR/${vm}-seed.iso"
  log "booting $vm  mac=$mac  ($why)"
  sudo qemu-img create -q -f qcow2 -F qcow2 -b "$BASE_IMAGE" "$disk" "$VM_DISK" >/dev/null
  build_dhcp_seed "$seed" "$vm" "$SSH_PUB"
  sudo virt-install --connect "$LIBVIRT_URI" --name "$vm" \
    --memory "$VM_MEM" --vcpus "$VM_CPUS" \
    --disk "path=$disk,format=qcow2,bus=virtio" \
    --disk "path=$seed,device=cdrom,readonly=on" \
    --network "network=labsim-ovs,portgroup=vlan${TEST_VLAN},model=virtio,mac=$mac" \
    --os-variant alpinelinux3.18 --graphics none --noautoconsole --import >/dev/null \
    || die "virt-install failed for $vm"
done

log "waiting ${BOOT_WAIT}s for boot + DHCP..."
sleep "$BOOT_WAIT"

# --- verdict --------------------------------------------------------------
# The lease table on the router is the authority: it says what the server
# decided, independent of whether the guest brought the interface up cleanly.
leases="$(router '/opt/vyatta/bin/vyatta-op-cmd-wrapper show dhcp server leases')"
echo
echo "=== router lease table ==="
echo "$leases"
echo

reserved_ips="$(cd "$SCRIPT_DIR/../migration" && python3 unifi-to-vyos.py --mode sim 2>/dev/null \
  | awk '/static-mapping .* ip-address/ {print $NF}')"

pass=0; fail=0
printf '%-19s %-16s %-16s %s\n' "MAC" "EXPECTED" "GOT" "RESULT"
for c in "${CASES[@]}"; do
  IFS='|' read -r mac expected why <<<"$c"
  got="$(echo "$leases" | awk -v m="$mac" 'tolower($0) ~ tolower(m) {print $1; exit}')"
  got="${got:-<none>}"
  if [ "$expected" = "POOL" ]; then
    if [ "$got" = "<none>" ]; then
      result="FAIL (no lease at all)"
    elif echo "$reserved_ips" | grep -qx "$got"; then
      result="FAIL (got a RESERVED address)"
    else
      result="pass"
    fi
  else
    [ "$got" = "$expected" ] && result="pass" || result="FAIL"
  fi
  [ "$result" = "pass" ] && pass=$((pass + 1)) || fail=$((fail + 1))
  printf '%-19s %-16s %-16s %s\n' "$mac" "$expected" "$got" "$result"
  printf '  %s\n' "$why"
done

echo
log "$pass passed, $fail failed"
[ "$KEEP" -eq 1 ] && log "VMs left running (--keep). Remove with: $0 --clean" || cleanup_vms
[ "$fail" -eq 0 ] || exit 1