#!/usr/bin/env bash
# discover-fortigate.sh — pull DHCP lease list from a FortiGate.
#
# SSHes to a FortiGate admin account, runs `execute dhcp lease-list`,
# parses the output into TSV (IP, MAC, hostname, vdom, source).
#
# Usage:
#   scripts/discover-fortigate.sh <fortigate-host>
#
# Env overrides:
#   FORTIGATE_SSH_USER  default: admin
#
# Output: TSV on stdout, one lease per line:
#   IP  MAC  HOSTNAME  VDOM  SOURCE
# Where SOURCE is "fortigate:<host>" so multiple runs can be concatenated
# and still identified.
#
# Example:
#   scripts/discover-fortigate.sh ana-fw.phasefinal.com > leases-ana.tsv
#
# Requires: ssh config (or `~/.ssh/config` host alias) for the FortiGate,
# with key auth OR interactive password. FortiGate's SSH expects admin-level
# credentials.

set -euo pipefail

if [ -z "${1:-}" ]; then
  echo "usage: $(basename "$0") <fortigate-host>" >&2
  exit 2
fi

ARG="$1"
# Accept either "host" or "user@host" — don't double-prefix the user.
if [[ "$ARG" == *@* ]]; then
  TARGET="$ARG"
  HOST="${ARG##*@}"
else
  TARGET="${FORTIGATE_SSH_USER:-admin}@${ARG}"
  HOST="$ARG"
fi
USER="${FORTIGATE_SSH_USER:-admin}"

# FortiGate CLI command. `execute dhcp lease-list all` dumps every vdom.
# If the device is single-vdom, `execute dhcp lease-list` (no arg) also works.
#
# Do NOT suppress stderr — FortiGate's error messages are the main debugging
# signal when the command returns empty. Let them flow to the caller.

raw=""
# Try the plain form first — works across FortiOS versions. `all` is
# an interface name on some devices but errors with "Interface name
# 'all' does not exist" on others (including 7.x ana-gw). Fall back
# to `all` only if the plain form returns no lease rows.
for cmd in 'execute dhcp lease-list' 'execute dhcp lease-list all'; do
  out=$(ssh -o StrictHostKeyChecking=accept-new "$TARGET" "$cmd" || true)
  # Accept only output that actually contains at least one IP-like token.
  # This filters out "Interface name X does not exist" and similar noise
  # without needing to enumerate every FortiOS error phrase.
  if grep -qE '([0-9]+\.){3}[0-9]+' <<<"$out"; then
    raw="$out"
    break
  fi
done

if [ -z "$raw" ]; then
  echo "error: no lease data from $HOST" >&2
  echo "       — try: ssh $TARGET  then run:  execute dhcp lease-list" >&2
  exit 1
fi

# Parse.
#
# Real FortiOS output looks like:
#
#   ana-gw # internal
#     IP            MAC-Address        Hostname       VCI         SSID  AP  SERVER-ID  Expiry
#     10.250.0.105  90:5a:08:98:e2:29  ana-ml2                                4          Sun Apr 26 17:02:23 2026
#   mgmt
#     IP            MAC-Address        Hostname       …
#     10.250.250.50 7c:c2:55:60:fe:8a  ANA-ML2-BMC   udhcp 1.32.1            6          Tue Apr 28 18:56:08 2026
#
# Interfaces are flush-left (no indent); lease rows + column headers are
# indented. First line embeds the shell prompt ("ana-gw # ") followed by
# the first interface name. Hostnames don't contain spaces in practice.

awk -v src="fortigate:${HOST}" '
  # Flush-left non-blank line = interface name (sometimes preceded by prompt)
  /^[^[:space:]]/ && NF >= 1 {
    # "<prompt> # <iface>" form — take the last field
    if (NF >= 3 && $2 == "#") { iface = $NF; next }
    # plain "<iface>" form
    if (NF == 1)              { iface = $1; next }
    # prompt-with-trailing-% or similar closing lines — skip
    next
  }
  # Column header row
  /^[[:space:]]+IP[[:space:]]/ { next }
  # Blank / whitespace-only
  /^[[:space:]]*$/             { next }
  # Lease row
  $1 ~ /^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$/ && $2 ~ /^[0-9a-fA-F]{2}:/ {
    ip   = $1
    mac  = tolower($2)   # normalize to lowercase — matches UniFi output
    host = (NF >= 3 && $3 !~ /^[0-9]+$/) ? $3 : "-"
    printf "%s\t%s\t%s\t%s\t%s\n", ip, mac, host, (iface==""?"-":iface), src
  }
' <<<"$raw"