#!/usr/bin/env bash
# ─────────────────────────────────────────────────────────────────────
# Lab PXE Bastion — discover-first bare-metal provisioning
#
# Default mode: DISCOVER. Any machine that PXE boots gets inventoried
# and powered off. You review what appeared, then promote to install.
#
# Usage:
#   sudo bash bastion.sh                          # start bastion (discover mode)
#   bash bastion.sh install <mac> <hostname>       # queue discovered machine for install
#   bash bastion.sh list                           # show discovered/queued machines
#
# Flow:
#   1. Start bastion        →  sudo bash bastion.sh
#   2. Power on machine     →  PXE boots, hardware discovered, powers off
#   3. Queue for install    →  bash bastion.sh install aa:bb:cc:dd:ee:ff puppet
#   4. Power on again       →  PXE boots, Fedora installed, reboots into OS
#
# Requirements: Fedora/RHEL host with dnsmasq, python3, curl
# ─────────────────────────────────────────────────────────────────────
set -euo pipefail

# ──── Configuration (override via environment) ────────────────────
FEDORA_VERSION="${FEDORA_VERSION:-43}"
ARCH="${ARCH:-x86_64}"
HTTP_PORT="${HTTP_PORT:-8080}"
TIMEZONE="${TIMEZONE:-Europe/London}"
LOCALE="${LOCALE:-en_GB.UTF-8}"
BASTION_DIR="${BASTION_DIR:-/tmp/lab-bastion}"

# ──── Colors ──────────────────────────────────────────────────────
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'
CYAN='\033[0;36m'; BOLD='\033[1m'; NC='\033[0m'

log()  { echo -e "${GREEN}[bastion]${NC} $*"; }
warn() { echo -e "${YELLOW}[bastion]${NC} $*"; }
err()  { echo -e "${RED}[bastion]${NC} $*" >&2; }
die()  { err "$@"; exit 1; }

# ──── Subcommand handling ─────────────────────────────────────────
CMD="${1:-serve}"

case "$CMD" in
    install)
        [[ $# -ge 3 ]] || { echo "Usage: bastion.sh install <mac> <hostname> [--disk <dev>]"; exit 1; }
        MAC="$2"
        HOSTNAME="$3"
        DISK="${5:-}"  # --disk <dev>
        PAYLOAD="{\"mac\":\"$MAC\",\"hostname\":\"$HOSTNAME\""
        [[ -n "$DISK" ]] && PAYLOAD="$PAYLOAD,\"disk\":\"$DISK\""
        PAYLOAD="$PAYLOAD}"
        RESULT=$(curl -sf -X POST "http://localhost:${HTTP_PORT}/api/install" \
            -H "Content-Type: application/json" \
            -d "$PAYLOAD" 2>&1) || die "Cannot reach bastion at localhost:${HTTP_PORT}. Is it running?"
        echo "$RESULT" | python3 -m json.tool 2>/dev/null || echo "$RESULT"
        echo ""
        echo "Power on the machine to start Fedora installation."
        exit 0
        ;;
    list)
        RESULT=$(curl -sf "http://localhost:${HTTP_PORT}/api/machines" 2>&1) || \
            die "Cannot reach bastion at localhost:${HTTP_PORT}. Is it running?"
        echo "$RESULT" | python3 -c "
import sys, json
state = json.load(sys.stdin)

discovered = state.get('discovered', {})
queue = state.get('install_queue', {})
installed = state.get('installed', {})

print()
print('\033[1mDISCOVERED\033[0m')
if discovered:
    print(f'  {\"MAC\":<20} {\"CPU\":<32} {\"CORES\":<6} {\"RAM\":<8} {\"ARCH\":<10} {\"PRODUCT\"}')
    for mac, hw in discovered.items():
        status = ' [QUEUED]' if mac in queue else ''
        print(f'  {mac:<20} {hw.get(\"cpu_model\",\"?\"):<32} {hw.get(\"cpu_cores\",\"?\"):<6} {str(hw.get(\"memory_gb\",\"?\"))+\"GB\":<8} {hw.get(\"arch\",\"?\"):<10} {hw.get(\"product\",\"?\")}{status}')
else:
    print('  (none — PXE boot a machine to discover it)')

print()
print('\033[1mINSTALL QUEUE\033[0m')
if queue:
    for mac, cfg in queue.items():
        print(f'  {mac:<20} → hostname={cfg.get(\"hostname\",\"?\")}')
else:
    print('  (none)')

print()
print('\033[1mINSTALLED\033[0m')
if installed:
    for mac, info in installed.items():
        print(f'  {mac:<20} → {info.get(\"hostname\",\"?\")}  ({info.get(\"installed_at\",\"?\")})')
else:
    print('  (none)')
print()
" 2>/dev/null || echo "$RESULT"
        exit 0
        ;;
    serve) ;;  # continue below
    *)
        echo "Usage: bastion.sh [serve|install <mac> <hostname>|list]"
        exit 1
        ;;
esac

# ══════════════════════════════════════════════════════════════════
# SERVE MODE — start the bastion
# ══════════════════════════════════════════════════════════════════

# ──── Preflight ───────────────────────────────────────────────────
[[ $EUID -eq 0 ]] || die "Must run as root (need DHCP/TFTP ports). Use: sudo bash bastion.sh"

command -v python3 >/dev/null || die "python3 not found"
command -v curl    >/dev/null || die "curl not found"

if ! command -v dnsmasq >/dev/null; then
    log "Installing dnsmasq..."
    if command -v dnf >/dev/null; then
        dnf install -y dnsmasq
    elif command -v apt-get >/dev/null; then
        apt-get install -y dnsmasq
    else
        die "Cannot install dnsmasq — install it manually"
    fi
fi

# ──── Auto-detect network ────────────────────────────────────────
IFACE="${IFACE:-$(ip route | awk '/default/ {print $5; exit}')}"
SERVER_IP="$(ip -4 addr show "$IFACE" | awk '/inet / {split($2,a,"/"); print a[1]; exit}')"
NETWORK="$(echo "$SERVER_IP" | awk -F. '{print $1"."$2"."$3".0"}')"

[[ -n "$SERVER_IP" ]] || die "Cannot detect IP on interface $IFACE"
log "Interface: ${BOLD}$IFACE${NC}  IP: ${BOLD}$SERVER_IP${NC}  Network: ${BOLD}$NETWORK${NC}"

# ──── Auto-detect SSH pubkey ──────────────────────────────────────
SSH_PUBKEY="${SSH_PUBKEY:-}"
if [[ -z "$SSH_PUBKEY" ]]; then
    REAL_HOME="${HOME}"
    [[ -n "${SUDO_USER:-}" ]] && REAL_HOME="$(getent passwd "$SUDO_USER" | cut -d: -f6)"
    for keyfile in "$REAL_HOME/.ssh/id_ed25519.pub" "$REAL_HOME/.ssh/id_rsa.pub" "$REAL_HOME/.ssh/id_ecdsa.pub"; do
        [[ -f "$keyfile" ]] && { SSH_PUBKEY="$keyfile"; break; }
    done
fi

SSH_KEY_CONTENT=""
if [[ -n "$SSH_PUBKEY" && -f "$SSH_PUBKEY" ]]; then
    SSH_KEY_CONTENT="$(cat "$SSH_PUBKEY")"
    log "SSH key: ${BOLD}$SSH_PUBKEY${NC}"
else
    warn "No SSH public key found. Set SSH_PUBKEY=/path/to/key.pub"
    warn "Install mode will use root password 'changeme' as fallback."
fi

# ──── Prepare directories ────────────────────────────────────────
TFTPDIR="$BASTION_DIR/tftp"
HTTPDIR="$BASTION_DIR/http"
STATEFILE="$BASTION_DIR/state.json"
mkdir -p "$TFTPDIR" "$HTTPDIR"

# Initialize state if not present
[[ -f "$STATEFILE" ]] || echo '{"discovered":{},"install_queue":{},"installed":{}}' > "$STATEFILE"

# ──── Cleanup handler ─────────────────────────────────────────────
DNSMASQ_PID=""
HTTP_PID=""
FW_OPENED=false

cleanup() {
    echo ""
    log "Shutting down..."
    [[ -n "$HTTP_PID" ]]    && kill "$HTTP_PID"    2>/dev/null && log "Stopped HTTP server"
    [[ -n "$DNSMASQ_PID" ]] && kill "$DNSMASQ_PID" 2>/dev/null && log "Stopped dnsmasq"

    if $FW_OPENED && command -v firewall-cmd >/dev/null; then
        log "Removing firewall rules..."
        firewall-cmd --quiet --remove-service=dhcp       2>/dev/null || true
        firewall-cmd --quiet --remove-service=tftp       2>/dev/null || true
        firewall-cmd --quiet --remove-port=${HTTP_PORT}/tcp 2>/dev/null || true
        firewall-cmd --quiet --remove-port=4011/udp      2>/dev/null || true
    fi

    log "State preserved in $STATEFILE"
    log "Restart bastion with: sudo bash bastion.sh"
}
trap cleanup EXIT INT TERM

# ──── Download artifacts (cached) ─────────────────────────────────
download() {
    local url="$1" dest="$2" label="$3"
    if [[ -f "$dest" ]]; then
        log "  ${label} — cached"
        return
    fi
    log "  ${label} — downloading..."
    curl -# -L -o "$dest" "$url" || die "Failed to download $label from $url"
}

FEDORA_MIRROR="https://download.fedoraproject.org/pub/fedora/linux/releases/${FEDORA_VERSION}/Everything/${ARCH}/os"

log "Fetching boot artifacts (Fedora ${FEDORA_VERSION} ${ARCH})..."
download "https://boot.ipxe.org/undionly.kpxe"   "$TFTPDIR/undionly.kpxe"   "iPXE BIOS"
download "https://boot.ipxe.org/ipxe.efi"        "$TFTPDIR/ipxe.efi"        "iPXE UEFI"
download "${FEDORA_MIRROR}/images/pxeboot/vmlinuz"    "$HTTPDIR/vmlinuz"     "Fedora kernel"
download "${FEDORA_MIRROR}/images/pxeboot/initrd.img" "$HTTPDIR/initrd.img"  "Fedora initrd"

# ──── Generate discovery kickstart ────────────────────────────────
# Boots Fedora installer env, collects hardware info, POSTs to bastion, powers off.
# Never touches the disk.
cat > "$HTTPDIR/discover.ks" << 'DISCOVER_KS'
# Lab Bastion — Discovery Mode
# Collects hardware inventory and powers off. Does NOT install anything.

%pre --erroronfail --log=/tmp/discover.log
#!/bin/bash
set -x

# ── Collect hardware info from /proc, /sys, and available tools ──

MAC=$(ip link show | awk '/ether/ && !/00:00:00:00/ {print $2; exit}')
PRODUCT=$(cat /sys/class/dmi/id/product_name 2>/dev/null || echo "unknown")
BOARD=$(cat /sys/class/dmi/id/board_name 2>/dev/null || echo "unknown")
SERIAL=$(cat /sys/class/dmi/id/product_serial 2>/dev/null || echo "unknown")
MANUFACTURER=$(cat /sys/class/dmi/id/sys_vendor 2>/dev/null || echo "unknown")
CPUMODEL=$(grep -m1 'model name' /proc/cpuinfo | cut -d: -f2 | sed 's/^ //')
CPUCORES=$(grep -c '^processor' /proc/cpuinfo)
MEMGB=$(awk '/MemTotal/ {printf "%d", $2/1024/1024}' /proc/meminfo)
ARCHTYPE=$(uname -m)

# Disk info — lsblk is available in Anaconda
DISKS_JSON=$(lsblk -Jb -o NAME,SIZE,TYPE,MODEL 2>/dev/null | python3 -c "
import sys, json
data = json.load(sys.stdin)
disks = [d for d in data.get('blockdevices', []) if d.get('type') == 'disk']
result = []
for d in disks:
    size_gb = round(int(d.get('size', 0)) / 1073741824, 1)
    result.append({
        'name': d.get('name', '?'),
        'size_gb': size_gb,
        'model': (d.get('model') or 'unknown').strip()
    })
print(json.dumps(result))
" 2>/dev/null || echo '[]')

# Network interfaces
NICS_JSON=$(ip -j link show 2>/dev/null | python3 -c "
import sys, json
nics = json.load(sys.stdin)
result = []
for n in nics:
    if n.get('link_type') == 'loopback':
        continue
    result.append({
        'name': n.get('ifname', '?'),
        'mac': n.get('address', '?'),
        'state': n.get('operstate', '?')
    })
print(json.dumps(result))
" 2>/dev/null || echo '[]')

# ── Build and POST discovery payload ──

PAYLOAD=$(python3 -c "
import json
print(json.dumps({
    'mac': '$MAC',
    'product': '$PRODUCT',
    'board': '$BOARD',
    'serial': '$SERIAL',
    'manufacturer': '$MANUFACTURER',
    'cpu_model': '$CPUMODEL',
    'cpu_cores': int('$CPUCORES' or 0),
    'memory_gb': int('$MEMGB' or 0),
    'arch': '$ARCHTYPE',
    'disks': $DISKS_JSON,
    'nics': $NICS_JSON
}))
")

# POST to bastion — try curl first, fall back to python3 urllib
BASTION_URL="__BASTION_URL__/api/discover"

if command -v curl >/dev/null 2>&1; then
    curl -sf -X POST "$BASTION_URL" \
        -H "Content-Type: application/json" \
        -d "$PAYLOAD" || true
else
    python3 -c "
import urllib.request
req = urllib.request.Request('$BASTION_URL',
    data=b'''$PAYLOAD''',
    headers={'Content-Type': 'application/json'})
try:
    urllib.request.urlopen(req, timeout=10)
except Exception as e:
    print(f'POST failed: {e}')
"
fi

# ── Power off — do NOT let Anaconda proceed ──
echo ""
echo "=== Discovery complete, powering off ==="
echo ""
sleep 3
echo 1 > /proc/sys/kernel/sysrq
echo o > /proc/sysrq-trigger
sleep 5
poweroff -f

%end

# Anaconda should never get here, but just in case:
poweroff
DISCOVER_KS

# Patch in the bastion URL
sed -i "s|__BASTION_URL__|http://${SERVER_IP}:${HTTP_PORT}|g" "$HTTPDIR/discover.ks"

# ──── Generate iPXE boot script ───────────────────────────────────
# Initial iPXE script chains to /dispatch with the MAC, so the server
# can route to discover or install mode per machine.
cat > "$HTTPDIR/boot.ipxe" << IPXE
#!ipxe

echo
echo ============================================
echo   Lab PXE Bastion
echo   Contacting server for instructions...
echo ============================================
echo

chain http://${SERVER_IP}:${HTTP_PORT}/dispatch?mac=\${net0/mac}
IPXE

# ──── Write the HTTP server ──────────────────────────────────────
cat > "$BASTION_DIR/server.py" << 'PYSERVER'
#!/usr/bin/env python3
"""Lab PXE Bastion — HTTP server with discovery API and per-MAC iPXE dispatch."""

import json
import os
import sys
import time
import fcntl
from http.server import HTTPServer, SimpleHTTPRequestHandler
from urllib.parse import urlparse, parse_qs
from datetime import datetime

# Config from argv
HTTP_DIR    = sys.argv[1]
STATE_FILE  = sys.argv[2]
SERVER_IP   = sys.argv[3]
HTTP_PORT   = int(sys.argv[4])
FEDORA_VER  = sys.argv[5]
FEDORA_MIRROR = sys.argv[6]
SSH_KEY     = sys.argv[7] if len(sys.argv) > 7 else ""
TIMEZONE    = sys.argv[8] if len(sys.argv) > 8 else "Europe/London"
LOCALE      = sys.argv[9] if len(sys.argv) > 9 else "en_GB.UTF-8"

# ── State management (file-backed, lock-protected) ───────────────

def load_state():
    try:
        with open(STATE_FILE) as f:
            return json.load(f)
    except (FileNotFoundError, json.JSONDecodeError):
        return {"discovered": {}, "install_queue": {}, "installed": {}}

def save_state(state):
    tmp = STATE_FILE + ".tmp"
    with open(tmp, 'w') as f:
        json.dump(state, f, indent=2)
    os.replace(tmp, STATE_FILE)

# ── Kickstart generation ─────────────────────────────────────────

def generate_kickstart(hostname, disk="", ssh_key=""):
    disk_cmds = "clearpart --all --initlabel\nautopart --type=plain"
    if disk:
        disk_cmds = f"ignoredisk --only-use={disk}\nclearpart --all --initlabel --drives={disk}\nautopart --type=plain"

    if ssh_key:
        auth = f'rootpw --lock\nsshkey --username=root "{ssh_key}"'
    else:
        auth = 'rootpw --plaintext changeme'

    return f"""# Lab Bastion — Fedora {FEDORA_VER} install
# Generated: {datetime.now().isoformat()}
# Target: {hostname}

text
reboot

lang {LOCALE}
keyboard uk
timezone {TIMEZONE} --utc

network --bootproto=dhcp --activate --hostname={hostname}

{auth}

{disk_cmds}

bootloader --append="console=tty0 console=ttyS0,115200n8"

url --mirrorlist=https://mirrors.fedoraproject.org/mirrorlist?repo=fedora-$releasever&arch=$basearch

%packages
@core
@server-product
openssh-server
vim-enhanced
tmux
git
curl
python3
lshw
dmidecode
dnf-plugins-core
%end

%post --log=/root/bastion-post-install.log
#!/bin/bash
set -x
systemctl enable --now sshd
sed -i 's/^#\\?PermitRootLogin.*/PermitRootLogin prohibit-password/' /etc/ssh/sshd_config
sed -i 's/^#\\?PasswordAuthentication.*/PasswordAuthentication no/' /etc/ssh/sshd_config
hostnamectl set-hostname {hostname}
echo "Provisioned by lab-bastion on $(date -Iseconds)" > /etc/lab-provisioned
echo "# Lab node — puppet enrollment pending" > /root/README
%end
"""

# ── Pretty terminal output ────────────────────────────────────────

RESET  = "\033[0m"
BOLD   = "\033[1m"
GREEN  = "\033[0;32m"
YELLOW = "\033[1;33m"
CYAN   = "\033[0;36m"
RED    = "\033[0;31m"

def print_discovery(mac, hw, is_new):
    """Print a discovered machine to the bastion terminal."""
    label = "NEW MACHINE DISCOVERED" if is_new else "MACHINE RE-DISCOVERED"
    color = GREEN if is_new else YELLOW

    # Format disk summary
    disks = hw.get('disks', [])
    disk_str = ", ".join(
        f"{d.get('size_gb', '?')}GB {d.get('model', '?')}"
        for d in disks
    ) or "none detected"

    # Format NIC summary
    nics = hw.get('nics', [])
    nic_str = ", ".join(n.get('name', '?') for n in nics) or "none"

    print(f"\n{color}{BOLD}{'═' * 60}")
    print(f"  {label}")
    print(f"{'═' * 60}{RESET}")
    print(f"  {BOLD}MAC:{RESET}      {mac}")
    print(f"  {BOLD}Product:{RESET}  {hw.get('manufacturer', '?')} {hw.get('product', '?')}")
    print(f"  {BOLD}CPU:{RESET}      {hw.get('cpu_model', '?')} ({hw.get('cpu_cores', '?')} cores)")
    print(f"  {BOLD}RAM:{RESET}      {hw.get('memory_gb', '?')} GB")
    print(f"  {BOLD}Arch:{RESET}     {hw.get('arch', '?')}")
    print(f"  {BOLD}Disks:{RESET}    {disk_str}")
    print(f"  {BOLD}NICs:{RESET}     {nic_str}")
    print(f"  {BOLD}Serial:{RESET}   {hw.get('serial', '?')}")
    print()
    print(f"  {CYAN}To install Fedora on this machine:{RESET}")
    print(f"  {BOLD}bash bastion.sh install {mac} <hostname>{RESET}")
    print(f"\n{'─' * 60}\n", flush=True)

def print_install_queued(mac, hostname):
    print(f"\n{GREEN}{BOLD}  INSTALL QUEUED{RESET}")
    print(f"  {mac} → hostname={BOLD}{hostname}{RESET}")
    print(f"  PXE boot the machine to start Fedora installation.")
    print(f"\n{'─' * 60}\n", flush=True)

def print_install_started(mac, hostname):
    print(f"\n{CYAN}{BOLD}  INSTALL STARTED{RESET}")
    print(f"  {mac} → {BOLD}{hostname}{RESET}")
    print(f"  Serving Fedora {FEDORA_VER} installer + kickstart...")
    print(f"\n{'─' * 60}\n", flush=True)

# ── HTTP Handler ──────────────────────────────────────────────────

class BastionHandler(SimpleHTTPRequestHandler):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, directory=HTTP_DIR, **kwargs)

    def log_message(self, format, *args):
        """Suppress default HTTP access logs — we have our own output."""
        pass

    def send_text(self, code, text, content_type="text/plain"):
        self.send_response(code)
        self.send_header("Content-Type", content_type)
        self.end_headers()
        self.wfile.write(text.encode())

    def send_json(self, code, data):
        self.send_text(code, json.dumps(data, indent=2), "application/json")

    def do_GET(self):
        parsed = urlparse(self.path)

        # ── iPXE dispatch: route to discover or install based on MAC ──
        if parsed.path == "/dispatch":
            params = parse_qs(parsed.query)
            mac = params.get("mac", [""])[0].lower().replace("-", ":")
            state = load_state()

            if mac in state.get("install_queue", {}):
                cfg = state["install_queue"][mac]
                hostname = cfg.get("hostname", "lab-node")
                print_install_started(mac, hostname)
                script = f"""#!ipxe

echo
echo =============================================
echo   Lab PXE Bastion — INSTALLING Fedora {FEDORA_VER}
echo   Target: {hostname}
echo   MAC:    {mac}
echo =============================================
echo

kernel http://{SERVER_IP}:{HTTP_PORT}/vmlinuz inst.ks=http://{SERVER_IP}:{HTTP_PORT}/ks?mac={mac} inst.repo={FEDORA_MIRROR} inst.text
initrd http://{SERVER_IP}:{HTTP_PORT}/initrd.img
boot
"""
                self.send_text(200, script)
            else:
                print(f"  {YELLOW}PXE request from {mac} → discovery mode{RESET}", flush=True)
                script = f"""#!ipxe

echo
echo =============================================
echo   Lab PXE Bastion — DISCOVERY MODE
echo   MAC: {mac}
echo   Collecting hardware info...
echo =============================================
echo

kernel http://{SERVER_IP}:{HTTP_PORT}/vmlinuz inst.ks=http://{SERVER_IP}:{HTTP_PORT}/discover.ks inst.text
initrd http://{SERVER_IP}:{HTTP_PORT}/initrd.img
boot
"""
                self.send_text(200, script)
            return

        # ── Per-MAC kickstart for install mode ──
        if parsed.path == "/ks":
            params = parse_qs(parsed.query)
            mac = params.get("mac", [""])[0].lower().replace("-", ":")
            state = load_state()
            cfg = state.get("install_queue", {}).get(mac, {})
            ks = generate_kickstart(
                hostname=cfg.get("hostname", "lab-node"),
                disk=cfg.get("disk", ""),
                ssh_key=SSH_KEY,
            )
            self.send_text(200, ks)
            return

        # ── API: list machines ──
        if parsed.path == "/api/machines":
            self.send_json(200, load_state())
            return

        # ── Static files (vmlinuz, initrd, discover.ks, etc.) ──
        super().do_GET()

    def do_POST(self):
        parsed = urlparse(self.path)
        content_length = int(self.headers.get("Content-Length", 0))
        body = self.rfile.read(content_length)

        # ── Discovery report from PXE-booted machine ──
        if parsed.path == "/api/discover":
            try:
                data = json.loads(body)
            except json.JSONDecodeError:
                self.send_json(400, {"error": "invalid JSON"})
                return

            mac = data.get("mac", "unknown").lower()
            data["last_seen"] = datetime.now().isoformat()

            state = load_state()
            is_new = mac not in state.get("discovered", {})
            if is_new:
                data["first_seen"] = data["last_seen"]
            else:
                data["first_seen"] = state["discovered"][mac].get("first_seen", data["last_seen"])

            state.setdefault("discovered", {})[mac] = data
            save_state(state)

            print_discovery(mac, data, is_new)

            self.send_json(200, {"status": "ok", "mac": mac, "new": is_new})
            return

        # ── Queue a machine for install ──
        if parsed.path == "/api/install":
            try:
                data = json.loads(body)
            except json.JSONDecodeError:
                self.send_json(400, {"error": "invalid JSON"})
                return

            mac = data.get("mac", "").lower().replace("-", ":")
            hostname = data.get("hostname", "lab-node")
            disk = data.get("disk", "")

            if not mac:
                self.send_json(400, {"error": "mac is required"})
                return

            state = load_state()
            state.setdefault("install_queue", {})[mac] = {
                "hostname": hostname,
                "disk": disk,
                "queued_at": datetime.now().isoformat(),
            }
            save_state(state)

            print_install_queued(mac, hostname)

            self.send_json(200, {
                "status": "queued",
                "mac": mac,
                "hostname": hostname,
                "message": "PXE boot the machine to start installation",
            })
            return

        self.send_json(404, {"error": "not found"})


def run_server():
    server = HTTPServer(("0.0.0.0", HTTP_PORT), BastionHandler)
    print(f"HTTP server listening on :{HTTP_PORT}", flush=True)
    server.serve_forever()


if __name__ == "__main__":
    run_server()
PYSERVER

# ──── Generate dnsmasq config ─────────────────────────────────────
cat > "$BASTION_DIR/dnsmasq.conf" << DNSMASQ
# Lab PXE Bastion — dnsmasq config
# ProxyDHCP mode: adds PXE options without replacing existing DHCP

# Disable DNS (we only want DHCP/TFTP)
port=0

# Listen on the right interface
interface=${IFACE}
bind-interfaces

# ProxyDHCP — works alongside existing DHCP (UniFi etc)
dhcp-range=${NETWORK},proxy

# TFTP for initial PXE boot
enable-tftp
tftp-root=${TFTPDIR}

# Detect client architecture
dhcp-match=set:bios,option:client-arch,0
dhcp-match=set:efi64,option:client-arch,7
dhcp-match=set:efi64,option:client-arch,9

# Detect iPXE clients (already chainloaded)
dhcp-userclass=set:ipxe,iPXE

# First PXE boot → serve iPXE binary via TFTP
dhcp-boot=tag:bios,tag:!ipxe,undionly.kpxe
dhcp-boot=tag:efi64,tag:!ipxe,ipxe.efi

# iPXE clients → chain to boot script via HTTP
dhcp-boot=tag:ipxe,http://${SERVER_IP}:${HTTP_PORT}/boot.ipxe

# Verbose logging
log-dhcp
DNSMASQ

# ──── Open firewall ──────────────────────────────────────────────
if command -v firewall-cmd >/dev/null && firewall-cmd --state >/dev/null 2>&1; then
    log "Opening firewall ports (DHCP, TFTP, HTTP:${HTTP_PORT})..."
    firewall-cmd --quiet --add-service=dhcp
    firewall-cmd --quiet --add-service=tftp
    firewall-cmd --quiet --add-port=${HTTP_PORT}/tcp
    firewall-cmd --quiet --add-port=4011/udp 2>/dev/null || true
    FW_OPENED=true
fi

# ──── Stop conflicting services ───────────────────────────────────
if systemctl is-active --quiet dnsmasq 2>/dev/null; then
    warn "System dnsmasq is running — stopping it temporarily"
    systemctl stop dnsmasq
fi

# ──── Start HTTP server ──────────────────────────────────────────
log "Starting HTTP server on :${HTTP_PORT}..."
python3 "$BASTION_DIR/server.py" \
    "$HTTPDIR" \
    "$STATEFILE" \
    "$SERVER_IP" \
    "$HTTP_PORT" \
    "$FEDORA_VERSION" \
    "$FEDORA_MIRROR" \
    "$SSH_KEY_CONTENT" \
    "$TIMEZONE" \
    "$LOCALE" &
HTTP_PID=$!
sleep 1

if ! kill -0 "$HTTP_PID" 2>/dev/null; then
    die "HTTP server failed to start — is port ${HTTP_PORT} in use?"
fi

# ──── Start dnsmasq ──────────────────────────────────────────────
log "Starting PXE server (proxyDHCP on ${IFACE})..."
echo ""
echo -e "${CYAN}${BOLD}════════════════════════════════════════════════════════════${NC}"
echo -e "${CYAN}${BOLD}  Lab PXE Bastion — Discovery Mode${NC}"
echo -e "${CYAN}${BOLD}════════════════════════════════════════════════════════════${NC}"
echo ""
echo -e "  Network:   ${BOLD}${NETWORK}/24${NC} via ${BOLD}${IFACE}${NC}"
echo -e "  HTTP:      ${BOLD}http://${SERVER_IP}:${HTTP_PORT}/${NC}"
echo -e "  OS:        ${BOLD}Fedora ${FEDORA_VERSION} (${ARCH})${NC}"
echo -e "  State:     ${BOLD}${STATEFILE}${NC}"
echo ""
echo -e "  ${YELLOW}PXE boot any machine on this network.${NC}"
echo -e "  ${YELLOW}It will be inventoried and powered off automatically.${NC}"
echo ""
echo -e "  Commands (from another terminal):"
echo -e "    ${BOLD}bash bastion.sh list${NC}                         — show machines"
echo -e "    ${BOLD}bash bastion.sh install <mac> <hostname>${NC}     — queue install"
echo ""
echo -e "  Press ${BOLD}Ctrl-C${NC} to stop."
echo ""
echo -e "${CYAN}──── Waiting for PXE boot requests... ────${NC}"
echo ""

dnsmasq --no-daemon --conf-file="$BASTION_DIR/dnsmasq.conf" &
DNSMASQ_PID=$!

wait "$DNSMASQ_PID" || {
    err "dnsmasq exited unexpectedly. Check if another DHCP/TFTP service is running."
    err "Try: ss -ulnp | grep -E ':(67|69|4011) '"
    exit 1
}