#!/bin/bash
# pxe-pick-nic - choose the imaging NIC by its cable, not by its MAC.
#
# Usage:
#   sudo pxe-pick-nic          interactive: pull the imaging switch cable and
#                              plug it back in; the port that blinked is it
#   pxe-pick-nic --auto        for the first-boot playbook: the one wired
#                              port with a link, if it is isolated
#   pxe-pick-nic --show        print the saved choice
#
# The choice is saved in /etc/pxe-server/nic-pick.conf (imaging_mac=...).
# The playbook uses it when site-vars.yml leaves pxe_nic_mac empty, and pins
# that MAC from then on: the box never re-guesses at boot, so a moved cable
# can never move 172.16.9.1 and DHCP onto another port.
#
# A port is accepted only when pxe-nic-check calls it "isolated": no GE
# switch (LLDP/CDP), no 802.1X port authentication, and no other DHCP
# server answering one DHCPDISCOVER. It listens CHECK_WINDOW seconds
# (default 65: a Cisco switch sends CDP every 60 s). --auto also refuses a
# port where more than ARP_MAX (default 10) hosts were heard: an imaging
# switch at install time is nearly silent, a GE access VLAN is not, and
# that holds even when the switch sends no LLDP/CDP and the probe is lost
# (STP still listening). Anything else is refused and nothing is saved.
# The check uses the site's corp_domain_suffixes and allowed_switches
# (PXE_PICK_SUFFIXES / PXE_PICK_ALLOWED from the playbook, else
# /etc/pxe-server/nic.conf), so a managed imaging switch can be allowed.
#
# Install rule for sites: during the install plug in ONLY the imaging
# switch. Plug the GE cable into any other port afterwards; the catch-all
# netplan entry makes it a DHCP client and dnsmasq never serves it.
#
# Exit: 0 chosen (or --show found one), 1 refused / could not choose,
# 2 usage.

set -u

CONF=${PXE_PICK_CONF:-/etc/pxe-server/nic-pick.conf}
HINT=${PXE_PICK_HINT:-/var/lib/pxe-firstboot/nic-hint}
SYSFS_NET=${SYSFS_NET:-/sys/class/net}
NIC_CHECK=${PXE_NIC_CHECK:-/usr/local/sbin/pxe-nic-check}
CHECK_WINDOW=${CHECK_WINDOW:-65}
ARP_MAX=${PXE_PICK_ARP_MAX:-10}
NIC_CONF=${PXE_NIC_CONF:-/etc/pxe-server/nic.conf}
SETTLE=${PXE_PICK_SETTLE:-8}
WAIT=${PXE_PICK_WAIT:-180}

log() { echo "pxe-pick-nic: $*"; }

# Physical wired Ethernet ports: a device behind them, ARP type 1, not
# wireless, not a bridge.
wired() {
    local d n
    for d in "$SYSFS_NET"/*; do
        n=${d##*/}
        [ -e "$d/device" ] || continue
        [ "$(cat "$d/type" 2>/dev/null)" = 1 ] || continue
        [ -e "$d/wireless" ] || [ -e "$d/phy80211" ] && continue
        [ -e "$d/bridge" ] && continue
        echo "$n"
    done
}

mac_of() { cat "$SYSFS_NET/$1/address" 2>/dev/null; }

# 1 = link, 0 = none. carrier is unreadable while the port is admin-down.
carrier() { [ "$(cat "$SYSFS_NET/$1/carrier" 2>/dev/null)" = 1 ] && echo 1 || echo 0; }

# Bring every wired port up (no address is added) so its link can be seen.
links_up() {
    local n up=0
    for n in $(wired); do
        if ! ip link show dev "$n" 2>/dev/null | grep -q '[<,]UP[,>]'; then
            ip link set dev "$n" up 2>/dev/null && up=1
        fi
    done
    [ "$up" -eq 1 ] && sleep "$SETTLE"
    return 0
}

table() {
    local n
    for n in $(wired); do
        printf '    %-16s %s  link %s\n' "$n" "$(mac_of "$n")" \
            "$([ "$(carrier "$n")" = 1 ] && echo up || echo down)"
    done
}

# A pxe-nic-check config with only the site settings that matter here
# (never the MAC keys: the port under test is given with --iface).
check_config() {
    local f=$1 suffixes=${PXE_PICK_SUFFIXES:-} allowed=${PXE_PICK_ALLOWED:-}
    if [ -z "$suffixes$allowed" ] && [ -f "$NIC_CONF" ]; then
        suffixes=$(sed -n 's/^[[:space:]]*corp_domain_suffixes[[:space:]]*=[[:space:]]*//p' "$NIC_CONF" | tail -1)
        allowed=$(sed -n 's/^[[:space:]]*allowed_switches[[:space:]]*=[[:space:]]*//p' "$NIC_CONF" | tail -1)
    fi
    {
        echo "[nic]"
        [ -n "$suffixes" ] && echo "corp_domain_suffixes = $suffixes"
        [ -n "$allowed" ] && echo "allowed_switches = $allowed"
    } >"$f"
}

# check IFACE [strict]: prints the verdict, returns 0 only for "isolated"
# (with strict, also only when at most ARP_MAX hosts were heard).
check() {
    local out verdict arp why cfg
    if [ ! -f "$NIC_CHECK" ]; then
        log "ERROR: $NIC_CHECK not found"
        return 1
    fi
    cfg=$(mktemp)
    check_config "$cfg"
    log "checking $1 is isolated: listening ${CHECK_WINDOW}s for GE switches and DHCP servers..."
    out=$(python3 "$NIC_CHECK" --iface "$1" --probe --window "$CHECK_WINDOW" --json --config "$cfg" 2>/dev/null)
    rm -f "$cfg"
    read -r verdict arp why < <(python3 -c '
import json, sys
try:
    e = json.load(sys.stdin)["interfaces"][0]
    ev = "; ".join(x for x in e["evidence"] if not x.endswith("ARP hosts seen"))
    print(e["verdict"], int(e.get("arp_hosts") or 0), ev or "-")
except Exception as ex:
    print("error", 0, "pxe-nic-check gave no usable result: %s" % ex)
' <<<"$out")
    if [ "$verdict" = isolated ]; then
        if [ "$arp" -gt "$ARP_MAX" ]; then
            log "$1: isolated by its switch and DHCP, but $arp hosts were heard (more than $ARP_MAX)"
            if [ "${2:-}" = strict ]; then
                log "REFUSED: a busy LAN is not an imaging switch at install time; it may be a GE port"
                log "with LLDP/CDP off. Pick it by its cable (sudo pxe-pick-nic) or set pxe_nic_mac."
                return 1
            fi
            log "WARNING: make sure this is the imaging switch, not a GE port with LLDP/CDP turned off."
        else
            log "$1: isolated (no GE switch, no 802.1X, no other DHCP server, $arp hosts heard)"
        fi
        return 0
    fi
    log "$1: verdict $verdict: $why"
    case "$verdict" in
        ge) log "That port is on the GE network (or a managed switch not in allowed_switches). It must never be the imaging port." ;;
        foreign-dhcp) log "Another DHCP server answers on that port (a second PXE server, or a routed LAN)." ;;
        no-link) log "That port lost its link during the check." ;;
    esac
    return 1
}

save() {
    local n=$1 mac
    mac=$(mac_of "$n")
    mkdir -p "$(dirname "$CONF")"
    {
        echo "# Written by pxe-pick-nic $(date -Is). The imaging NIC; the playbook"
        echo "# pins it when site-vars.yml leaves pxe_nic_mac empty."
        echo "imaging_mac=$mac"
        echo "picked_iface=$n"
    } >"$CONF.tmp" && mv "$CONF.tmp" "$CONF"
    rm -f "$HINT"
    log "saved: imaging NIC = $n ($mac) in $CONF"
}

hint() {
    mkdir -p "$(dirname "$HINT")" 2>/dev/null
    { printf '%s\n' "$*" >"$HINT"; } 2>/dev/null
}

saved_mac() {
    sed -n 's/^imaging_mac=\(\([0-9a-f]\{2\}:\)\{5\}[0-9a-f]\{2\}\)$/\1/p' "$CONF" 2>/dev/null | tail -1
}

# --- modes ---------------------------------------------------------------------

mode_show() {
    local mac
    mac=$(saved_mac)
    [ -n "$mac" ] || { log "no saved choice in $CONF"; return 1; }
    echo "$mac"
}

mode_auto() {
    local mac n with=() names
    mac=$(saved_mac)
    if [ -n "$mac" ]; then
        for n in $(wired); do
            if [ "$(mac_of "$n")" = "$mac" ]; then
                log "using the saved choice: $n ($mac)"
                echo "$mac"
                return 0
            fi
        done
        log "the saved imaging NIC $mac is not on this box any more"
        hint "The saved imaging NIC $mac is gone (dongle replaced?). Log in and run: sudo pxe-pick-nic"
        return 1
    fi
    links_up
    for n in $(wired); do
        [ "$(carrier "$n")" = 1 ] && with+=("$n")
    done
    log "wired ports:"
    table
    if [ ${#with[@]} -eq 0 ]; then
        log "REFUSED: no wired port has a link. Plug the imaging switch into the port to use."
        hint "Could not pick the imaging port: no port has a link. Plug in the imaging switch, then log in and run: sudo pxe-pick-nic"
        return 1
    fi
    if [ ${#with[@]} -gt 1 ]; then
        names="${with[*]}"
        log "REFUSED: ${#with[@]} ports have a link ($names). During the install only the imaging switch may be plugged in."
        hint "Could not pick the imaging port: ${#with[@]} ports have a link ($names). Log in and run: sudo pxe-pick-nic"
        return 1
    fi
    n=${with[0]}
    if ! check "$n" strict; then
        hint "Could not pick the imaging port: $n is not isolated (see /var/log/first-boot.log). Plug only the imaging switch in, then log in and run: sudo pxe-pick-nic"
        return 1
    fi
    save "$n" >&2
    mac_of "$n"
}

# Wait until exactly one port changes its link; print "NAME NEWSTATE".
wait_change() {
    local -A before=()
    local n t=0 changed=()
    for n in $(wired); do before[$n]=$(carrier "$n"); done
    while [ "$t" -lt $((WAIT * 2)) ]; do
        sleep 0.5
        t=$((t + 1))
        changed=()
        for n in "${!before[@]}"; do
            [ "$(carrier "$n")" != "${before[$n]}" ] && changed+=("$n")
        done
        if [ ${#changed[@]} -eq 1 ]; then
            sleep 1   # a flapping connector settles
            n=${changed[0]}
            echo "$n $(carrier "$n")"
            return 0
        fi
        if [ ${#changed[@]} -gt 1 ]; then
            log "more than one port changed (${changed[*]}); move one cable at a time" >&2
            return 1
        fi
    done
    log "nothing changed in ${WAIT}s" >&2
    return 1
}

mode_interactive() {
    local n state again ans
    [ -t 0 ] || { log "interactive mode needs a terminal (use --auto from scripts)"; return 2; }
    [ "$(id -u)" -eq 0 ] || { log "run it as root: sudo pxe-pick-nic"; return 2; }
    links_up
    echo
    echo "Pick the imaging NIC: the port cabled to the isolated imaging switch."
    echo "Wired ports now:"
    table
    echo
    echo "Pull the imaging switch cable out of this server now"
    echo "(or, if it is not plugged in yet, plug it in now). Waiting up to ${WAIT}s..."
    read -r n state < <(wait_change) || return 1
    [ -n "${n:-}" ] || return 1
    if [ "$state" = 0 ]; then
        echo "$n lost its link. Plug the cable back into the SAME port..."
        read -r again state < <(wait_change) || return 1
        if [ "${again:-}" != "$n" ] || [ "$state" != 1 ]; then
            log "expected $n to come back up, saw ${again:-nothing} $state; try again"
            return 1
        fi
    fi
    sleep "$SETTLE"
    echo "Imaging port: $n ($(mac_of "$n"))"
    check "$n" || { log "nothing saved"; return 1; }
    read -rp "Use $n ($(mac_of "$n")) as the imaging NIC? [y/N] " ans
    [[ "$ans" =~ ^[Yy]$ ]] || { log "nothing saved"; return 1; }
    save "$n"
    if [ -e /var/lib/pxe-firstboot.failed ] && [ -x /opt/first-boot.sh ]; then
        read -rp "Run the server setup again now? [Y/n] " ans
        if [[ ! "$ans" =~ ^[Nn]$ ]]; then
            rm -f /var/lib/pxe-firstboot/attempts
            /opt/first-boot.sh 2>&1 | tee -a /var/log/first-boot.log
            return "${PIPESTATUS[0]}"
        fi
        echo "Run it later with: sudo rm -f /var/lib/pxe-firstboot/attempts; sudo /opt/first-boot.sh"
    else
        echo "Saved. If the server is already set up, re-run the playbook from the"
        echo "installer USB to move the imaging config onto this port."
    fi
    return 0
}

case "${1:-}" in
    --auto) mode_auto ;;
    --show) mode_show ;;
    "") mode_interactive ;;
    -h|--help) sed -n '2,/^$/p' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;;
    *) echo "usage: $0 [--auto | --show]" >&2; exit 2 ;;
esac