#!/usr/bin/env bash
# ─────────────────────────────────────────────────────────────────────
# Lab PXE Bastion — ephemeral PXE server for bare-metal provisioning
#
# Turns this machine into a temporary PXE boot server.  Target machines
# on the same network can PXE boot and get Fedora installed automatically.
#
# Usage:
#   sudo bash bastion.sh                        # interactive, auto-detect everything
#   sudo TARGET_HOSTNAME=puppet SSH_PUBKEY=~/.ssh/id_ed25519.pub bash bastion.sh
#
# Requirements: Fedora/RHEL host with dnsmasq, python3, curl
# ─────────────────────────────────────────────────────────────────────
set -euo pipefail

# ──── Defaults (override via environment) ──────────────────────────
FEDORA_VERSION="${FEDORA_VERSION:-41}"
ARCH="${ARCH:-x86_64}"
HTTP_PORT="${HTTP_PORT:-8080}"
TARGET_HOSTNAME="${TARGET_HOSTNAME:-lab-node}"
TARGET_DISK="${TARGET_DISK:-}"          # empty = anaconda auto-picks
SSH_PUBKEY="${SSH_PUBKEY:-}"            # path to .pub file, auto-detected
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; }

# ──── 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"

# Install dnsmasq if missing
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 ───────────────────────────────────────
if [[ -z "$SSH_PUBKEY" ]]; then
    # When run via sudo, check the real user's home
    REAL_HOME="${HOME}"
    if [[ -n "${SUDO_USER:-}" ]]; then
        REAL_HOME="$(getent passwd "$SUDO_USER" | cut -d: -f6)"
    fi
    for keyfile in "$REAL_HOME/.ssh/id_ed25519.pub" "$REAL_HOME/.ssh/id_rsa.pub" "$REAL_HOME/.ssh/id_ecdsa.pub"; do
        if [[ -f "$keyfile" ]]; then
            SSH_PUBKEY="$keyfile"
            break
        fi
    done
fi

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. Root password will be set to 'changeme'."
    warn "Set SSH_PUBKEY=/path/to/key.pub to use key-based auth instead."
    SSH_KEY_CONTENT=""
fi

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

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

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

    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-service=proxy-dhcp 2>/dev/null || true
    fi

    log "Done. Bastion artifacts remain in $BASTION_DIR"
    log "Re-run this script to reprovision. Remove with: rm -rf $BASTION_DIR"
}
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 kickstart ──────────────────────────────────────────
log "Generating kickstart for ${BOLD}${TARGET_HOSTNAME}${NC}..."

# Disk config
if [[ -n "$TARGET_DISK" ]]; then
    DISK_CMDS="ignoredisk --only-use=${TARGET_DISK}
clearpart --all --initlabel --drives=${TARGET_DISK}
autopart --type=plain"
else
    DISK_CMDS="clearpart --all --initlabel
autopart --type=plain"
fi

# Auth config
if [[ -n "$SSH_KEY_CONTENT" ]]; then
    AUTH_CMDS="rootpw --lock
sshkey --username=root \"${SSH_KEY_CONTENT}\""
else
    AUTH_CMDS='rootpw --plaintext changeme'
fi

cat > "$HTTPDIR/ks.cfg" << KICKSTART
# Lab Bastion — Fedora ${FEDORA_VERSION} kickstart
# Generated: $(date -Iseconds)
# Target: ${TARGET_HOSTNAME}

# Install mode
text
reboot

# Locale
lang ${LOCALE}
keyboard uk
timezone ${TIMEZONE} --utc

# Network
network --bootproto=dhcp --activate --hostname=${TARGET_HOSTNAME}

# Auth
${AUTH_CMDS}

# Disk
${DISK_CMDS}

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

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

# Packages — minimal server + essentials
%packages
@core
@server-product
openssh-server
vim-enhanced
tmux
git
curl
python3
dnf-plugins-core
%end

# Post-install
%post --log=/root/bastion-post-install.log
#!/bin/bash
set -x

# Ensure SSH is enabled
systemctl enable --now sshd

# Allow root SSH with key (password auth disabled)
sed -i 's/^#\?PermitRootLogin.*/PermitRootLogin prohibit-password/' /etc/ssh/sshd_config
sed -i 's/^#\?PasswordAuthentication.*/PasswordAuthentication no/' /etc/ssh/sshd_config

# Set hostname
hostnamectl set-hostname ${TARGET_HOSTNAME}

# Leave a breadcrumb
echo "Provisioned by lab-bastion on $(date -Iseconds)" > /etc/lab-provisioned

# Placeholder: puppet enrollment will go here later
# puppet is not installed yet — this IS the puppet server
echo "# Lab bootstrap node — puppet server setup pending" > /root/README

%end
KICKSTART

log "Kickstart written to ${HTTPDIR}/ks.cfg"

# ──── Generate iPXE boot script ───────────────────────────────────
cat > "$HTTPDIR/boot.ipxe" << IPXE
#!ipxe

echo
echo =======================================
echo   Lab PXE Bastion — Fedora ${FEDORA_VERSION}
echo   Target: ${TARGET_HOSTNAME}
echo =======================================
echo

kernel http://${SERVER_IP}:${HTTP_PORT}/vmlinuz inst.ks=http://${SERVER_IP}:${HTTP_PORT}/ks.cfg inst.repo=${FEDORA_MIRROR} inst.text
initrd http://${SERVER_IP}:${HTTP_PORT}/initrd.img
boot
IPXE

# ──── 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 (see what's happening)
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
    # ProxyDHCP uses port 4011
    firewall-cmd --quiet --add-port=4011/udp 2>/dev/null || true
    FW_OPENED=true
fi

# ──── Stop conflicting services ───────────────────────────────────
# dnsmasq might be running as a system service
if systemctl is-active --quiet dnsmasq 2>/dev/null; then
    warn "System dnsmasq is running — stopping it temporarily"
    systemctl stop dnsmasq
    RESTART_DNSMASQ=true
fi

# ──── Start HTTP server ───────────────────────────────────────────
log "Starting HTTP server on :${HTTP_PORT}..."
(cd "$HTTPDIR" && python3 -m http.server "$HTTP_PORT" --bind 0.0.0.0 >/dev/null 2>&1) &
HTTP_PID=$!
sleep 0.5

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

# ──── Start dnsmasq (proxyDHCP + TFTP) ────────────────────────────
log "Starting PXE server (proxyDHCP on ${IFACE})..."
echo ""
echo -e "${CYAN}${BOLD}════════════════════════════════════════════════════════${NC}"
echo -e "${CYAN}${BOLD}  PXE Bastion ready!${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 "  Hostname:   ${BOLD}${TARGET_HOSTNAME}${NC}"
echo -e "  Kickstart:  ${BOLD}http://${SERVER_IP}:${HTTP_PORT}/ks.cfg${NC}"
echo ""
echo -e "  ${YELLOW}Now PXE-boot the target machine.${NC}"
echo -e "  ${YELLOW}Set boot order to Network/PXE in BIOS, or use one-time boot menu.${NC}"
echo ""
echo -e "  Press ${BOLD}Ctrl-C${NC} to stop the bastion."
echo ""
echo -e "${CYAN}──── dnsmasq log (watch for DHCP/PXE requests) ────${NC}"
echo ""

# Run dnsmasq in foreground so logs stream to terminal
dnsmasq --no-daemon --conf-file="$BASTION_DIR/dnsmasq.conf" &
DNSMASQ_PID=$!

# Wait for dnsmasq — if it exits, something went wrong
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
}