#!/bin/bash
# pxe-netplan-apply.sh - install the PXE server's netplan config without
# ever leaving the box unreachable on the corporate side.
#
# Usage:
#   pxe-netplan-apply.sh STAGED_YAML PXE_NAME PXE_MAC [CORP_NAME CORP_MAC]
#   pxe-netplan-apply.sh --preflight STAGED_YAML PXE_NAME PXE_MAC [CORP_NAME CORP_MAC]
#   pxe-netplan-apply.sh --revert BACKUP_DIR      (the auto-revert timer)
#
# Without CORP_NAME/CORP_MAC (corp_nic_mac empty in the playbook) only the
# imaging NIC is pinned and renamed; the other wired ports are catch-all
# DHCP clients that keep their names. Every pre-flight refusal about the
# imaging NIC applies the same.
#
# Env (set by pxe_server_setup.yml):
#   PXE_SUBNET      imaging LAN, e.g. 172.16.9.0/24 (SSH peers inside it may
#                   legitimately come in on the imaging NIC)
#   PXE_SSH_CLIENT  IP of the admin's SSH client ($SSH_CONNECTION field 1);
#                   every established sshd peer is checked as well
#   PXE_NETPLAN_REVERT_AFTER  seconds before the auto-revert fires (120)
#   PXE_NETPLAN_CORP_WAIT     seconds to wait for the corp IPv4 again (30)
#
# Steps (apply):
#   1. validate STAGED_YAML in a scratch root (netplan generate --root-dir);
#      nothing on the box changes if it is invalid
#   2. stop here (UNCHANGED) if it is already installed, no other file
#      configures wired NICs, and both NICs carry their stable names
#   3. pre-flight: REFUSE (exit 1, nothing changed) when the NIC that
#      PXE_MAC names holds the default route, a DHCP lease, or the route to
#      an SSH client outside PXE_SUBNET. That NIC becomes static with no
#      gateway, so any of these means it is the corp-cabled one (swapped
#      MACs) and applying would cut the admin off.
#   4. back up /etc/netplan to /var/backups/netplan-<timestamp>/, record the
#      current NIC names, and arm an auto-revert (systemd-run timer) that
#      restores the backup unless step 8 cancels it
#   5. install it (other wired configs moved aside), netplan generate
#   6. rename the NICs to their stable names (the .link files netplan
#      writes do it on every later boot)
#   7. netplan apply
#   8. check: both names present, the imaging NIC has its IPv4, the corp
#      NIC has IPv4 again if it had one before (PXE_NETPLAN_CORP_WAIT s),
#      SSH clients still route out of the corp side; then cancel the revert
# Any failure in 5-8 restores the backup AND the previous NIC names, then
# exits 1. Prints CHANGED / UNCHANGED (or NEEDED for --preflight) as the
# last line.

set -u

NETPLAN_DIR=${NETPLAN_DIR:-/etc/netplan}
SYSFS_NET=${SYSFS_NET:-/sys/class/net}
BACKUP_ROOT=${BACKUP_ROOT:-/var/backups}
LEASE_DIR=${LEASE_DIR:-/run/systemd/netif/leases}
LOCK_FILE=${LOCK_FILE:-/run/pxe-netplan-apply.lock}
REVERT_AFTER=${PXE_NETPLAN_REVERT_AFTER:-120}
CORP_WAIT=${PXE_NETPLAN_CORP_WAIT:-30}
PXE_WAIT=${PXE_NETPLAN_PXE_WAIT:-20}
TARGET=$NETPLAN_DIR/50-cloud-init.yaml
SELF=$(readlink -f "$0")

log() { echo "pxe-netplan-apply: $*"; }

# --- NIC helpers -------------------------------------------------------------

# Name of the physical NIC that has MAC $1 (empty if none).
iface_of() {
    local d
    for d in "$SYSFS_NET"/*; do
        [ -e "$d/device" ] || continue
        if [ "$(cat "$d/address" 2>/dev/null)" = "$1" ]; then
            basename "$d"
            return 0
        fi
    done
    return 0
}

has_v4() {
    ip -4 -o addr show dev "$1" scope global 2>/dev/null | grep -q ' inet '
}

default_route_devs() {
    { ip -4 route show default; ip -6 route show default; } 2>/dev/null |
        awk '{for (i = 1; i < NF; i++) if ($i == "dev") print $(i + 1)}' | sort -u
}

# DHCP-assigned address (networkd and NetworkManager both flag those
# "dynamic"), or a networkd lease file for the link.
has_dhcp_lease() {
    local idx
    ip -o addr show dev "$1" 2>/dev/null | grep -qw dynamic && return 0
    idx=$(cat "$SYSFS_NET/$1/ifindex" 2>/dev/null)
    [ -n "$idx" ] && [ -s "$LEASE_DIR/$idx" ] && return 0
    return 1
}

route_dev() {
    ip -o route get "$1" 2>/dev/null |
        awk '{for (i = 1; i < NF; i++) if ($i == "dev") {print $(i + 1); exit}}'
}

ip2int() {
    local a b c d
    [[ "$1" =~ ^([0-9]+)\.([0-9]+)\.([0-9]+)\.([0-9]+)$ ]] || return 1
    a=${BASH_REMATCH[1]} b=${BASH_REMATCH[2]} c=${BASH_REMATCH[3]} d=${BASH_REMATCH[4]}
    echo $(( (a << 24) | (b << 16) | (c << 8) | d ))
}

in_subnet() {
    local ipn netn len mask
    [ -n "${2:-}" ] || return 1
    ipn=$(ip2int "$1") || return 1
    netn=$(ip2int "${2%/*}") || return 1
    len=${2#*/}
    [[ "$len" =~ ^[0-9]+$ ]] && [ "$len" -le 32 ] || return 1
    mask=$(( len == 0 ? 0 : (0xFFFFFFFF << (32 - len)) & 0xFFFFFFFF ))
    [ $(( ipn & mask )) -eq $(( netn & mask )) ]
}

# Admin SSH clients: PXE_SSH_CLIENT, our own SSH_CONNECTION, and every
# established peer of sshd (port 22). IPv4-mapped IPv6 is unwrapped.
ssh_peers() {
    {
        [ -n "${PXE_SSH_CLIENT:-}" ] && echo "$PXE_SSH_CLIENT"
        [ -n "${SSH_CONNECTION:-}" ] && echo "${SSH_CONNECTION%% *}"
        ss -Htn state established '( sport = :22 )' 2>/dev/null | awk '{print $NF}' |
            sed -E 's/:[0-9]+$//; s/^\[//; s/\]$//'
    } | sed -E 's/^::ffff://' | grep -E '^[0-9a-fA-F.:]+$' | sort -u
}

# --- pre-flight --------------------------------------------------------------

# 0 = safe to apply. Prints every reason it is not.
preflight() {
    local pif cif peer dev bad=0
    pif=$(iface_of "$PXE_MAC")
    if [ -z "$pif" ]; then
        log "REFUSED: no physical NIC has pxe_nic_mac $PXE_MAC"
        return 1
    fi
    if [ -n "$CORP_MAC" ]; then
        cif=$(iface_of "$CORP_MAC")
        if [ -z "$cif" ]; then
            log "REFUSED: no physical NIC has corp_nic_mac $CORP_MAC"
            return 1
        fi
    fi
    if default_route_devs | grep -qxF "$pif"; then
        log "REFUSED: $pif (pxe_nic_mac $PXE_MAC) holds the default route"
        bad=1
    fi
    if has_dhcp_lease "$pif"; then
        log "REFUSED: $pif (pxe_nic_mac $PXE_MAC) has a DHCP lease"
        bad=1
    fi
    while IFS= read -r peer; do
        [ -n "$peer" ] || continue
        dev=$(route_dev "$peer")
        [ "$dev" = "$pif" ] || continue
        if in_subnet "$peer" "${PXE_SUBNET:-}"; then
            continue    # an admin on the imaging LAN keeps 172.16.9.1
        fi
        log "REFUSED: the SSH client $peer is reached through $pif (pxe_nic_mac $PXE_MAC)"
        bad=1
    done < <(ssh_peers)
    if [ "$bad" -ne 0 ]; then
        log "The NIC named by pxe_nic_mac looks like the GE corporate one: the MACs are"
        log "probably swapped, or the cables are. It would become"
        log "static ${PXE_SUBNET:-imaging} with no gateway and cut this box off. Nothing was changed."
        log "NICs now: $(for d in "$SYSFS_NET"/*; do [ -e "$d/device" ] && printf '%s=%s ' "$(basename "$d")" "$(cat "$d/address" 2>/dev/null)"; done)"
        return 1
    fi
    return 0
}

# --- renaming ----------------------------------------------------------------

link_rename() {
    local cur=$1 new=$2
    log "renaming $cur to $new"
    if ip link set dev "$cur" down &&
       ip link set dev "$cur" name "$new" &&
       ip link set dev "$new" up; then
        return 0
    fi
    log "ERROR: rename $cur -> $new failed"
    ip link set dev "$cur" up 2>/dev/null
    ip link set dev "$new" up 2>/dev/null
    return 1
}

# set_names MAC1 NAME1 [MAC2 NAME2]: give each MAC its name. A NIC that must
# change goes through a temporary name first, so a swap (pxe0 <-> corp0)
# works in both directions. Temporary names carry this run's PID, so they
# never collide with a name left over from an earlier (failed) pass. A pair
# with an empty MAC or name is skipped (no pinned GE NIC).
set_names() {
    local -a macs=() want=() cur=()
    local i
    while [ $# -ge 2 ]; do
        if [ -n "$1" ] && [ -n "$2" ]; then
            macs+=("$1"); want+=("$2")
        fi
        shift 2
    done
    RENAME_SEQ=$(( ${RENAME_SEQ:-0} + 1 ))
    for i in "${!macs[@]}"; do
        cur[i]=$(iface_of "${macs[i]}")
        if [ -z "${cur[i]}" ]; then
            log "ERROR: no NIC with MAC ${macs[i]}"
            return 1
        fi
    done
    for i in "${!macs[@]}"; do
        [ "${cur[i]}" = "${want[i]}" ] && continue
        link_rename "${cur[i]}" "pxt$$-$RENAME_SEQ$i" || return 1
        cur[i]=pxt$$-$RENAME_SEQ$i
    done
    for i in "${!macs[@]}"; do
        [ "${cur[i]}" = "${want[i]}" ] && continue
        link_rename "${cur[i]}" "${want[i]}" || return 1
    done
    for i in "${!macs[@]}"; do
        if [ "$(iface_of "${macs[i]}")" != "${want[i]}" ]; then
            log "ERROR: MAC ${macs[i]} is not named ${want[i]} after renaming"
            return 1
        fi
    done
    return 0
}

names_ok() {
    [ "$(iface_of "$PXE_MAC")" = "$PXE_NAME" ] || return 1
    if [ -n "$CORP_MAC" ]; then
        [ "$(iface_of "$CORP_MAC")" = "$CORP_NAME" ]
    else
        [ -z "$(stray_names)" ]
    fi
}

# No pinned GE NIC: "MAC NAME" for every other physical NIC whose name the
# catch-all (en*, eth*) would not match - the previous imaging NIC still
# called pxe0 after a re-pick, or corp0 from a pinned config. Each gets the
# kernel-style MAC name enx<mac>, so it becomes a catch-all DHCP client and
# pxe0 is free for the new imaging NIC.
stray_names() {
    local d n m
    [ -z "$CORP_MAC" ] || return 0
    for d in "$SYSFS_NET"/*; do
        [ -e "$d/device" ] || continue
        n=${d##*/}
        m=$(cat "$d/address" 2>/dev/null)
        [ "$m" = "$PXE_MAC" ] && continue
        [ "$(cat "$d/type" 2>/dev/null)" = 1 ] || continue
        [ -e "$d/wireless" ] || [ -e "$d/phy80211" ] || [ -e "$d/bridge" ] && continue
        case "$n" in en*|eth*) continue ;; esac
        echo "$m enx${m//:/}"
    done
    return 0
}

# MACs of the physical NICs other than the imaging one that hold a DHCP
# (dynamic) IPv4: the GE side. They must have one again after the apply,
# or an admin on the GE side may have been cut off. Static addresses (the
# previous imaging NIC's 172.16.9.1) do not count. By MAC: names change.
other_v4_macs() {
    local d n m
    for d in "$SYSFS_NET"/*; do
        [ -e "$d/device" ] || continue
        n=${d##*/}
        m=$(cat "$d/address" 2>/dev/null)
        [ "$m" = "$PXE_MAC" ] && continue
        ip -4 -o addr show dev "$n" scope global dynamic 2>/dev/null | grep -q ' inet ' && echo "$m"
    done
    return 0
}

# --- backup / restore --------------------------------------------------------

# Other netplan files that configure wired NICs (installer leftovers, an
# en* DHCP match): they would put a second config on the two NICs. Files
# without "ethernets" (WiFi, NetworkManager wifi profiles) are left alone.
other_files() {
    find "$NETPLAN_DIR" -maxdepth 1 -type f \( -name '*.yaml' -o -name '*.yml' \) \
        ! -path "$TARGET" -exec grep -l 'ethernets' {} + 2>/dev/null | sort
}

# State of one apply, next to (not inside) the backup, so restoring the
# backup never copies it into /etc/netplan.
write_state() {
    {
        echo "NETPLAN_DIR=$NETPLAN_DIR"
        echo "PXE_MAC=$PXE_MAC"
        echo "ORIG_PXE_IF=$ORIG_PXE_IF"
        echo "CORP_MAC=$CORP_MAC"
        echo "ORIG_CORP_IF=$ORIG_CORP_IF"
        echo "EXTRA=$EXTRA"
    } >"$BACKUP.state"
}

read_state() {
    local k v
    while IFS='=' read -r k v; do
        case "$k" in
            NETPLAN_DIR) [[ "$v" =~ ^/[A-Za-z0-9/._-]*$ ]] && NETPLAN_DIR=$v ;;
            PXE_MAC|CORP_MAC) [[ "$v" =~ ^([0-9a-f]{2}:){5}[0-9a-f]{2}$ ]] && printf -v "$k" '%s' "$v" ;;
            ORIG_PXE_IF|ORIG_CORP_IF) [[ "$v" =~ ^[A-Za-z0-9._-]{0,15}$ ]] && printf -v "$k" '%s' "$v" ;;
            EXTRA) [[ "$v" =~ ^(([0-9a-f]{2}:){5}[0-9a-f]{2}=[A-Za-z0-9._-]{1,15},?)*$ ]] && EXTRA=$v ;;
        esac
    done <"$BACKUP.state"
    TARGET=$NETPLAN_DIR/50-cloud-init.yaml
}

cancel_revert() {
    [ -n "${REVERT_UNIT:-}" ] || return 0
    systemctl stop "$REVERT_UNIT.timer" "$REVERT_UNIT.service" >/dev/null 2>&1
    if systemctl is-active --quiet "$REVERT_UNIT.timer" 2>/dev/null; then
        log "WARNING: could not cancel $REVERT_UNIT.timer (the committed marker still stops it)"
    fi
}

# Put the backed-up config AND the previous NIC names back.
restore() {
    local rc=0
    log "restoring the previous config from $BACKUP"
    find "$NETPLAN_DIR" -maxdepth 1 -type f \( -name '*.yaml' -o -name '*.yml' \) -delete
    cp -a "$BACKUP/." "$NETPLAN_DIR/"
    if [ -n "${ORIG_PXE_IF:-}" ]; then
        local -a back=()
        local pair
        for pair in ${EXTRA//,/ }; do back+=("${pair%%=*}" "${pair#*=}"); done
        set_names "$PXE_MAC" "$ORIG_PXE_IF" "$CORP_MAC" "${ORIG_CORP_IF:-}" "${back[@]}" ||
            { log "ERROR: could not put the previous NIC names back"; rc=1; }
    fi
    netplan generate && netplan apply || { log "ERROR: netplan apply of the restored config failed"; rc=1; }
    touch "$BACKUP.reverted"
    return $rc
}

fail() {
    log "ERROR: $*"
    restore
    cancel_revert
    log "the previous config is back in place; nothing applied"
    exit 1
}

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

MODE=apply
case "${1:-}" in
    --preflight) MODE=preflight; shift ;;
    --revert)
        [ $# -eq 2 ] || { echo "usage: $0 --revert BACKUP_DIR" >&2; exit 2; }
        BACKUP=${2%/}
        [ -d "$BACKUP" ] && [ -f "$BACKUP.state" ] || { log "revert: no backup/state at $BACKUP"; exit 1; }
        exec 9>"$LOCK_FILE"
        flock -w 600 9 || log "revert: lock still held after 600 s; reverting anyway"
        if [ -e "$BACKUP.committed" ] || [ -e "$BACKUP.reverted" ]; then
            log "revert: $BACKUP already committed or reverted; nothing to do"
            exit 0
        fi
        PXE_MAC="" CORP_MAC="" ORIG_PXE_IF="" ORIG_CORP_IF="" EXTRA=""
        read_state
        log "revert: the apply was not confirmed within its window"
        restore
        exit $?
        ;;
esac

if [ $# -ne 3 ] && [ $# -ne 5 ]; then
    echo "usage: $0 [--preflight] STAGED_YAML PXE_NAME PXE_MAC [CORP_NAME CORP_MAC]" >&2
    exit 2
fi
STAGED=$1
PXE_NAME=$2
PXE_MAC=${3,,}
CORP_NAME=${4:-}
CORP_MAC=${5:-}
CORP_MAC=${CORP_MAC,,}

# 1. validate in a scratch root
SCRATCH=$(mktemp -d)
trap 'rm -rf "$SCRATCH"' EXIT
mkdir -p "$SCRATCH/etc/netplan"
install -m 600 "$STAGED" "$SCRATCH/etc/netplan/50-cloud-init.yaml"
if ! netplan generate --root-dir "$SCRATCH"; then
    log "ERROR: $STAGED does not validate; nothing changed"
    exit 1
fi

# 2. already in place?
if cmp -s "$STAGED" "$TARGET" && [ -z "$(other_files)" ] && names_ok; then
    log "netplan already in place, $PXE_NAME${CORP_NAME:+ and $CORP_NAME} present"
    echo UNCHANGED
    exit 0
fi

# 3. pre-flight
if ! preflight; then
    exit 1
fi
if [ "$MODE" = preflight ]; then
    log "pre-flight passed; the netplan config would change"
    echo NEEDED
    exit 0
fi

exec 9>"$LOCK_FILE"
flock -w 60 9 || { log "ERROR: another pxe-netplan-apply holds $LOCK_FILE; nothing changed"; exit 1; }

# 4. back up, record names, arm the auto-revert
ORIG_PXE_IF=$(iface_of "$PXE_MAC")
ORIG_CORP_IF=""
[ -n "$CORP_MAC" ] && ORIG_CORP_IF=$(iface_of "$CORP_MAC")
CORP_HAD_V4=0
[ -n "$ORIG_CORP_IF" ] && has_v4 "$ORIG_CORP_IF" && CORP_HAD_V4=1
OTHERS_V4=""
[ -z "$CORP_MAC" ] && OTHERS_V4=$(other_v4_macs)
# NICs to move off a name the catch-all would not match, recorded with
# their current names for the revert ("mac=name,...").
EXTRA=""
FREE=()
while read -r m newname; do
    [ -n "$m" ] || continue
    EXTRA="${EXTRA:+$EXTRA,}$m=$(iface_of "$m")"
    FREE+=("$m" "$newname")
done < <(stray_names)
BACKUP=$BACKUP_ROOT/netplan-$(date +%Y%m%d-%H%M%S)
mkdir -p "$BACKUP"
cp -a "$NETPLAN_DIR/." "$BACKUP/"
write_state
log "previous config saved in $BACKUP (NICs: $ORIG_PXE_IF${ORIG_CORP_IF:+, $ORIG_CORP_IF})"

REVERT_UNIT=""
if [ "$REVERT_AFTER" -gt 0 ] 2>/dev/null && command -v systemd-run >/dev/null; then
    REVERT_UNIT=pxe-netplan-revert-$(basename "$BACKUP" | tr -cd 'A-Za-z0-9-')
    if systemd-run --quiet --unit="$REVERT_UNIT" --on-active="$REVERT_AFTER" \
        --timer-property=AccuracySec=1s "$SELF" --revert "$BACKUP"; then
        log "auto-revert armed: $REVERT_UNIT.timer restores $BACKUP in ${REVERT_AFTER}s unless this run confirms"
    else
        log "WARNING: could not arm the auto-revert timer; relying on the in-script checks"
        REVERT_UNIT=""
    fi
else
    log "WARNING: auto-revert off (no systemd-run or PXE_NETPLAN_REVERT_AFTER=0)"
fi

# An SSH drop (HUP) must not stop us half-way; TERM/INT restore.
trap '' HUP
trap 'fail "interrupted"' INT TERM

# 5. install, moving aside the other wired configs (an en* DHCP match would
#    put a DHCP client on the imaging NIC too)
while IFS= read -r f; do
    [ -n "$f" ] || continue
    log "moving aside $f (copy in $BACKUP)"
    rm -f "$f"
done < <(other_files)
install -m 600 "$STAGED" "$TARGET"
netplan generate || fail "netplan generate failed"

# 6. stable names now, so dnsmasq/UFW/Samba (which use them) work before
#    the next reboot. Does not depend on the netplan version renaming.
[ ${#FREE[@]} -gt 0 ] && log "freeing names for the catch-all: $EXTRA"
set_names "$PXE_MAC" "$PXE_NAME" "$CORP_MAC" "$CORP_NAME" "${FREE[@]}" || fail "NIC rename failed"

# 7. apply
netplan apply || fail "netplan apply failed"

# 8. check
names_ok || fail "$PXE_NAME${CORP_NAME:+ / $CORP_NAME} not on MAC $PXE_MAC${CORP_MAC:+ / $CORP_MAC} after apply"
ok=0
for _ in $(seq 1 "$PXE_WAIT"); do
    if has_v4 "$PXE_NAME"; then ok=1; break; fi
    sleep 1
done
[ "$ok" -eq 1 ] || fail "$PXE_NAME has no IPv4 address after ${PXE_WAIT}s"
for m in $OTHERS_V4; do
    ok=0
    for _ in $(seq 1 "$CORP_WAIT"); do
        n=$(iface_of "$m")
        if [ -n "$n" ] && has_v4 "$n"; then ok=1; break; fi
        sleep 1
    done
    [ "$ok" -eq 1 ] || fail "the NIC with MAC $m had a DHCP IPv4 before and has none after ${CORP_WAIT}s"
done
if [ -z "$CORP_MAC" ]; then
    :
elif [ "$CORP_HAD_V4" -eq 1 ]; then
    ok=0
    for _ in $(seq 1 "$CORP_WAIT"); do
        if has_v4 "$CORP_NAME"; then ok=1; break; fi
        sleep 1
    done
    [ "$ok" -eq 1 ] || fail "$CORP_NAME had IPv4 before and has none after ${CORP_WAIT}s"
elif ! has_v4 "$CORP_NAME"; then
    log "WARNING: $CORP_NAME has no IPv4 address (it had none before either: corp cable unplugged or no DHCP)"
fi
while IFS= read -r peer; do
    [ -n "$peer" ] || continue
    dev=$(route_dev "$peer")
    if [ "$dev" = "$PXE_NAME" ] && ! in_subnet "$peer" "${PXE_SUBNET:-}"; then
        fail "the SSH client $peer now routes through $PXE_NAME"
    fi
    if [ -z "$dev" ]; then
        fail "no route to the SSH client $peer after apply"
    fi
done < <(ssh_peers)

touch "$BACKUP.committed"
cancel_revert
trap - INT TERM
ip -br addr show dev "$PXE_NAME" 2>/dev/null
[ -n "$CORP_NAME" ] && ip -br addr show dev "$CORP_NAME" 2>/dev/null
echo CHANGED
exit 0