#!/usr/bin/env python3
"""Commissioning agent: runs once inside the ephemeral live image.

Collects a hardware inventory, configures a BMC account via in-band ipmitool
(falls back to a server-provided out-of-band hint for virtual machines),
posts everything to rollcall and powers the machine off.

Kernel cmdline options:
  rollcall.server=http://10.213.77.1   (required) rollcall base URL
  rollcall.nopoweroff                  stay up (debugging)
  rollcall.bmcuser=commissioner        BMC account name to create
  rollcall.token=...                   agent token (sent as X-Rollcall-Token)
  rollcall.mode=deploy rollcall.job=.. run rollcall-deploy.py instead (OS deployment / release job)
"""
import glob
import json
import os
import re
import secrets
import socket
import string
import subprocess
import sys
import time
import urllib.error
import urllib.request

LOG = open("/dev/console", "w", buffering=1) if os.path.exists("/dev/console") else sys.stdout


def log(msg):
    line = f"[commission] {msg}"
    print(line, file=LOG)
    if LOG is not sys.stdout:
        print(line, flush=True)


def run(cmd, timeout=120, text=True):
    """Return (rc, stdout). Never raises."""
    try:
        p = subprocess.run(cmd, shell=isinstance(cmd, str), capture_output=True,
                           text=text, timeout=timeout)
        return p.returncode, p.stdout if text else p.stdout
    except FileNotFoundError:
        return 127, ""
    except subprocess.TimeoutExpired:
        return 124, ""


def out(cmd, timeout=120):
    rc, o = run(cmd, timeout)
    return o.strip() if rc == 0 else ""


def jout(cmd, timeout=120):
    o = out(cmd, timeout)
    try:
        return json.loads(o) if o else None
    except json.JSONDecodeError:
        return None


def readf(path, default=""):
    try:
        with open(path) as f:
            return f.read().strip()
    except OSError:
        return default


def cmdline():
    opts = {}
    for tok in readf("/proc/cmdline").split():
        if tok.startswith("rollcall."):
            k, _, v = tok[len("rollcall."):].partition("=")
            opts[k] = v if v else True
    return opts


# --------------------------------------------------------------------------- DMI
def dmi_system():
    keys = ["system-manufacturer", "system-product-name", "system-version",
            "system-serial-number", "system-uuid", "system-sku-number", "system-family",
            "baseboard-manufacturer", "baseboard-product-name", "baseboard-serial-number",
            "chassis-manufacturer", "chassis-type", "chassis-serial-number", "chassis-asset-tag",
            "bios-vendor", "bios-version", "bios-release-date"]
    d = {k: out(["dmidecode", "-s", k]) for k in keys}
    # fallbacks via sysfs
    d["system-uuid"] = d["system-uuid"] or readf("/sys/class/dmi/id/product_uuid")
    d["system-serial-number"] = d["system-serial-number"] or readf("/sys/class/dmi/id/product_serial")
    return d


def dmi_blocks(dtype):
    """Parse `dmidecode -t N` into a list of dicts (one per handle)."""
    txt = out(["dmidecode", "-t", str(dtype)])
    blocks, cur = [], None
    for line in txt.splitlines():
        if line.startswith("Handle "):
            cur = {"_handle": line.split(",")[0].split()[1]}
            blocks.append(cur)
        elif cur is not None and line.startswith("\t") and ":" in line and not line.startswith("\t\t"):
            k, _, v = line.strip().partition(":")
            cur[k.strip()] = v.strip()
    return blocks


def memory():
    dimms = []
    for b in dmi_blocks(17):
        size = b.get("Size", "")
        if not size or size.startswith("No Module"):
            dimms.append({"locator": b.get("Locator"), "bank": b.get("Bank Locator"), "populated": False})
            continue
        dimms.append({
            "locator": b.get("Locator"), "bank": b.get("Bank Locator"), "populated": True,
            "size": size, "type": b.get("Type"), "type_detail": b.get("Type Detail"),
            "speed": b.get("Speed"), "configured_speed": b.get("Configured Memory Speed"),
            "manufacturer": b.get("Manufacturer"), "part_number": b.get("Part Number"),
            "serial": b.get("Serial Number"), "rank": b.get("Rank"), "form_factor": b.get("Form Factor"),
        })
    total_kb = 0
    m = re.search(r"MemTotal:\s+(\d+)", readf("/proc/meminfo"))
    if m:
        total_kb = int(m.group(1))
    return {"total_bytes": total_kb * 1024, "dimm_slots": len(dimms),
            "dimms_populated": sum(1 for d in dimms if d["populated"]), "dimms": dimms}


def cpu():
    j = jout(["lscpu", "-J"]) or {}
    flat = {}
    for e in j.get("lscpu", []):
        flat[e["field"].rstrip(":")] = e["data"]
        for c in e.get("children", []) or []:
            flat[c["field"].rstrip(":")] = c["data"]
    want = ["Architecture", "CPU(s)", "Thread(s) per core", "Core(s) per socket", "Socket(s)",
            "NUMA node(s)", "Vendor ID", "Model name", "CPU family", "Model", "Stepping",
            "CPU max MHz", "CPU min MHz", "L1d cache", "L2 cache", "L3 cache", "Virtualization", "Flags"]
    res = {k: flat.get(k) for k in want if k in flat}
    flags = res.pop("Flags", "") or ""
    res["flags_count"] = len(flags.split())
    res["features"] = sorted(f for f in flags.split() if f in
                             ("vmx", "svm", "avx", "avx2", "avx512f", "sse4_2", "aes", "sha_ni", "amx_tile"))
    return res


# --------------------------------------------------------------------------- storage
def storage():
    lsblk = jout(["lsblk", "-J", "-O", "-b", "-d"]) or {}
    disks = []
    for d in lsblk.get("blockdevices", []):
        if d.get("type") not in ("disk", "nvme") and not d.get("name", "").startswith("nvme"):
            continue
        if d.get("name", "").startswith(("loop", "ram", "sr", "zram")):
            continue
        disks.append({k: d.get(k) for k in ("name", "size", "model", "serial", "vendor", "rev", "tran",
                                             "rota", "wwn", "hctl", "subsystems", "log-sec", "phy-sec",
                                             "rm", "ro", "path", "pttype", "fstype")})
    nvme = jout(["nvme", "list", "-o", "json"]) or {}
    smart = {}
    for d in disks:
        s = jout(["smartctl", "-i", "-H", "-j", d["path"]], timeout=30) or {}
        if s:
            smart[d["name"]] = {
                "model_family": s.get("model_family"), "firmware": s.get("firmware_version"),
                "rotation_rate": s.get("rotation_rate"), "form_factor": (s.get("form_factor") or {}).get("name"),
                "healthy": (s.get("smart_status") or {}).get("passed"),
                "power_on_hours": (s.get("power_on_time") or {}).get("hours"),
                "interface_speed": ((s.get("interface_speed") or {}).get("current") or {}).get("string"),
            }
    ctrls = [p for p in pci_devices() if p["class_id"][:2] == "01"]  # mass storage controllers
    return {"disks": disks, "disk_count": len(disks),
            "total_bytes": sum(int(d["size"] or 0) for d in disks),
            "nvme": nvme.get("Devices"), "smart": smart, "controllers": ctrls}


# --------------------------------------------------------------------------- PCI
_PCI_CACHE = None


def pci_devices():
    global _PCI_CACHE
    if _PCI_CACHE is not None:
        return _PCI_CACHE
    txt = out(["lspci", "-Dvmmnnk"])
    devs, cur = [], {}
    for line in txt.splitlines() + [""]:
        if not line.strip():
            if cur:
                devs.append(cur)
            cur = {}
            continue
        k, _, v = line.partition(":")
        cur[k.strip()] = v.strip()
    res = []
    for d in devs:
        m_cls = re.search(r"\[([0-9a-f]{4})\]$", d.get("Class", ""))
        m_ven = re.search(r"\[([0-9a-f]{4})\]$", d.get("Vendor", ""))
        m_dev = re.search(r"\[([0-9a-f]{4})\]$", d.get("Device", ""))
        cls = m_cls.group(1) if m_cls else "0000"
        kind = {"01": "storage", "02": "network", "03": "gpu", "0c": "serial-bus", "06": "bridge"}.get(cls[:2], "other")
        if cls == "0302" or (cls == "1200"):
            kind = "accelerator"
        res.append({
            "slot": d.get("Slot"), "class": re.sub(r"\s*\[[0-9a-f]{4}\]$", "", d.get("Class", "")),
            "class_id": cls, "kind": kind,
            "vendor": re.sub(r"\s*\[[0-9a-f]{4}\]$", "", d.get("Vendor", "")),
            "device": re.sub(r"\s*\[[0-9a-f]{4}\]$", "", d.get("Device", "")),
            "vendor_id": m_ven.group(1) if m_ven else None, "device_id": m_dev.group(1) if m_dev else None,
            "subsystem": d.get("SDevice"), "driver": d.get("Driver"), "rev": d.get("Rev"),
        })
    _PCI_CACHE = res
    return res


def gpus():
    return [p for p in pci_devices() if p["kind"] in ("gpu", "accelerator")]


# --------------------------------------------------------------------------- network
def ethtool_info(ifn):
    info = {}
    for line in out(["ethtool", "-i", ifn]).splitlines():
        k, _, v = line.partition(":")
        if k in ("driver", "version", "firmware-version", "bus-info"):
            info[k.replace("-", "_")] = v.strip()
    et = out(["ethtool", ifn])
    m = re.search(r"Speed:\s*(\S+)", et)
    info["speed"] = m.group(1) if m else None
    m = re.search(r"Duplex:\s*(\S+)", et)
    info["duplex"] = m.group(1) if m else None
    m = re.search(r"Link detected:\s*(\S+)", et)
    info["link"] = (m.group(1) == "yes") if m else None
    m = re.search(r"Supported ports:\s*\[([^\]]*)\]", et)
    info["ports"] = m.group(1).split() if m else None
    sup = re.search(r"Supported link modes:(.*?)Supported pause", et, re.S)
    if sup:
        info["supported_link_modes"] = sup.group(1).split()
    pm = out(["ethtool", "-P", ifn])
    m = re.search(r"([0-9a-f]{2}(?::[0-9a-f]{2}){5})", pm)
    info["permanent_mac"] = m.group(1) if m else None
    return info


def network():
    links = jout(["ip", "-j", "-d", "link"]) or []
    addrs = {a["ifname"]: a for a in (jout(["ip", "-j", "addr"]) or [])}
    nics = []
    for l in links:
        ifn = l["ifname"]
        if ifn == "lo" or l.get("link_type") not in ("ether", "infiniband"):
            continue
        sysd = f"/sys/class/net/{ifn}"
        if not os.path.exists(f"{sysd}/device"):
            continue  # virtual (bridge, bond, vlan, veth ...)
        pci = os.path.basename(os.path.realpath(f"{sysd}/device"))
        nic = {
            "name": ifn, "mac": l.get("address"), "mtu": l.get("mtu"), "operstate": l.get("operstate"),
            "link_type": l.get("link_type"), "pci": pci,
            "sriov_totalvfs": readf(f"{sysd}/device/sriov_totalvfs") or None,
            "numa_node": readf(f"{sysd}/device/numa_node") or None,
            "ipv4": [a["local"] + "/" + str(a["prefixlen"]) for a in addrs.get(ifn, {}).get("addr_info", [])
                     if a.get("family") == "inet"],
            "ipv6": [a["local"] + "/" + str(a["prefixlen"]) for a in addrs.get(ifn, {}).get("addr_info", [])
                     if a.get("family") == "inet6" and a.get("scope") == "global"],
        }
        if l.get("link_type") == "ether":
            nic.update(ethtool_info(ifn))
        nics.append(nic)
    return nics


def infiniband():
    hcas = []
    for d in sorted(glob.glob("/sys/class/infiniband/*")):
        name = os.path.basename(d)
        ports = []
        for p in sorted(glob.glob(f"{d}/ports/*")):
            ports.append({
                "port": int(os.path.basename(p)), "state": readf(f"{p}/state"),
                "phys_state": readf(f"{p}/phys_state"), "rate": readf(f"{p}/rate"),
                "link_layer": readf(f"{p}/link_layer"), "lid": readf(f"{p}/lid"),
                "sm_lid": readf(f"{p}/sm_lid"), "port_guid": readf(f"{p}/gids/0") or None,
            })
        hcas.append({
            "name": name, "node_guid": readf(f"{d}/node_guid"), "sys_image_guid": readf(f"{d}/sys_image_guid"),
            "fw_ver": readf(f"{d}/fw_ver"), "hca_type": readf(f"{d}/hca_type"), "board_id": readf(f"{d}/board_id"),
            "node_desc": readf(f"{d}/node_desc"),
            "pci": os.path.basename(os.path.realpath(f"{d}/device")) if os.path.exists(f"{d}/device") else None,
            "ports": ports,
        })
    return {"hcas": hcas, "ibstat": out(["ibstat"], timeout=30) or None}


def lldp(wait=25):
    if run(["which", "lldpcli"])[0] != 0:
        return None
    deadline = time.time() + wait
    neigh = None
    while time.time() < deadline:
        j = jout(["lldpcli", "-f", "json", "show", "neighbors", "details"], timeout=20)
        if j and j.get("lldp") and j["lldp"].get("interface"):
            neigh = j["lldp"]["interface"]
            break
        time.sleep(3)
    return neigh


# --------------------------------------------------------------------------- BMC
def gen_password(n=16):
    alphabet = string.ascii_letters + string.digits
    return "".join(secrets.choice(alphabet) for _ in range(n))


def ipmi(args, timeout=30):
    return run(["ipmitool"] + args, timeout=timeout)


def parse_kv(txt):
    d = {}
    for line in txt.splitlines():
        if ":" in line:
            k, _, v = line.partition(":")
            k = k.strip()
            if k and k not in d:
                d[k] = v.strip()
    return d


def bmc_inband(user_name):
    """Discover the BMC in-band, create/refresh user_name with ADMINISTRATOR on the LAN channel."""
    for mod in ("ipmi_msghandler", "ipmi_devintf", "ipmi_si", "ipmi_ssif"):
        run(["modprobe", mod])
    time.sleep(2)
    if not os.path.exists("/dev/ipmi0"):
        return None
    res = {"source": "inband", "errors": []}
    rc, mc = ipmi(["mc", "info"])
    res["mc_info"] = parse_kv(mc) if rc == 0 else None
    rc, fru = ipmi(["fru", "print", "0"], timeout=60)
    res["fru"] = parse_kv(fru) if rc == 0 else None
    # find LAN channel
    chan, lan = None, None
    for c in range(1, 15):
        rc, o = ipmi(["lan", "print", str(c)])
        if rc == 0 and "IP Address" in o:
            kv = parse_kv(o)
            if kv.get("IP Address") and kv["IP Address"] != "0.0.0.0":
                chan, lan = c, kv
                break
            chan, lan = chan or c, lan or kv
    if chan is None:
        res["errors"].append("no LAN channel found")
        return res
    res["channel"] = chan
    res["lan"] = {k: lan.get(k) for k in ("IP Address Source", "IP Address", "Subnet Mask", "MAC Address",
                                          "Default Gateway IP", "802.1q VLAN ID", "Cipher Suite Priv Max")}
    res["address"] = lan.get("IP Address")
    res["mac"] = lan.get("MAC Address")
    # user table
    rc, ul = ipmi(["user", "list", str(chan)])
    users = []
    for line in ul.splitlines()[1:]:
        m = re.match(r"^\s*(\d+)\s+(.*?)\s+(true|false)\s+(true|false)\s+(true|false)\s+(\S.*)$", line)
        if m:
            users.append({"id": int(m.group(1)), "name": m.group(2).strip(), "callin": m.group(3) == "true",
                          "link": m.group(4) == "true", "ipmi_msg": m.group(5) == "true", "priv": m.group(6).strip()})
    res["users_before"] = users
    uid = next((u["id"] for u in users if u["name"] == user_name), None)
    if uid is None:
        uid = next((u["id"] for u in users if u["id"] >= 3 and not u["name"]), None)
    if uid is None:
        res["errors"].append("no free user slot")
        return res
    pw = gen_password()
    steps = [
        ["user", "set", "name", str(uid), user_name],
        ["user", "set", "password", str(uid), pw],
        ["user", "priv", str(uid), "4", str(chan)],
        ["channel", "setaccess", str(chan), str(uid), "callin=on", "ipmi=on", "link=on", "privilege=4"],
        ["user", "enable", str(uid)],
    ]
    ok = True
    for s in steps:
        rc, o = ipmi(s)
        if rc != 0:
            ok = False
            res["errors"].append(f"ipmitool {' '.join(s[:3])} rc={rc}")
    # some BMCs (e.g. iDRAC) also want SOL payload access enabled
    ipmi(["sol", "payload", "enable", str(chan), str(uid)])
    res.update({"user": user_name, "user_id": uid, "password": pw if ok else None, "configured": ok})
    if res["address"] and res["address"] != "0.0.0.0":
        rc, _ = ipmi(["-I", "lanplus", "-H", res["address"], "-U", user_name, "-P", pw, "mc", "info"], timeout=20)
        res["verified_oob"] = (rc == 0)
    return res


AGENT_TOKEN = ""


def agent_headers():
    h = {"Content-Type": "application/json"}
    if AGENT_TOKEN:
        h["X-Rollcall-Token"] = AGENT_TOKEN
    return h


def bmc_hint(server, macs):
    try:
        req = urllib.request.Request(f"{server}/api/v1/bmc-hint?macs={','.join(macs)}", headers=agent_headers())
        with urllib.request.urlopen(req, timeout=10) as r:
            h = json.load(r)
    except (urllib.error.URLError, json.JSONDecodeError, OSError):
        return None
    if not h or not h.get("address"):
        return None
    res = {"source": "hint", "address": h["address"], "port": h.get("port", 623),
           "user": h.get("user"), "password": h.get("password"), "errors": []}
    base = ["-I", "lanplus", "-H", h["address"], "-p", str(h.get("port", 623)), "-U", h.get("user", ""),
            "-P", h.get("password", "")]
    rc, mc = ipmi(base + ["mc", "info"], timeout=20)
    res["mc_info"] = parse_kv(mc) if rc == 0 else None
    res["verified_oob"] = (rc == 0)
    res["configured"] = (rc == 0)
    if rc != 0:
        res["errors"].append("hinted BMC did not answer")
    return res


# --------------------------------------------------------------------------- main
def post(server, payload, tries=30):
    data = json.dumps(payload).encode()
    req = urllib.request.Request(f"{server}/api/v1/commission", data=data, headers=agent_headers(), method="POST")
    for i in range(tries):
        try:
            with urllib.request.urlopen(req, timeout=30) as r:
                return json.load(r)
        except (urllib.error.URLError, OSError) as e:
            log(f"post failed ({e}), retry {i + 1}/{tries}")
            time.sleep(5)
    return None


def main():
    opts = cmdline()
    if opts.get("mode") == "deploy":
        # same image, different job: hand over to the deploy/release agent
        os.execv("/usr/bin/python3", ["python3", "/usr/local/bin/rollcall-deploy.py"])
    server = opts.get("server")
    if not server:
        log("rollcall.server= missing on kernel cmdline; nothing to do")
        return 2
    server = server.rstrip("/")
    global AGENT_TOKEN
    AGENT_TOKEN = opts.get("token") if isinstance(opts.get("token"), str) else ""
    user_name = opts.get("bmcuser") or "commissioner"
    started = time.time()
    log(f"starting, server={server}")

    inv = {"agent_version": "1", "started_at": started, "hostname": socket.gethostname(),
           "kernel": readf("/proc/version"), "errors": []}
    for key, fn in [("system", dmi_system), ("cpu", cpu), ("memory", memory), ("storage", storage),
                    ("network", network), ("infiniband", infiniband), ("pci", pci_devices), ("gpus", gpus)]:
        try:
            inv[key] = fn()
            log(f"collected {key}")
        except Exception as e:  # noqa: BLE001
            inv[key] = None
            inv["errors"].append(f"{key}: {e!r}")
            log(f"FAILED {key}: {e!r}")
    inv["tpm"] = [os.path.basename(p) for p in glob.glob("/sys/class/tpm/tpm*")]
    inv["lldp"] = lldp()
    log("collected lldp")

    macs = [n["mac"] for n in (inv.get("network") or []) if n.get("mac")]
    try:
        bmc = bmc_inband(user_name)
        if bmc is None:
            log("no /dev/ipmi0, asking server for an out-of-band hint")
            bmc = bmc_hint(server, macs)
    except Exception as e:  # noqa: BLE001
        bmc = {"source": "error", "errors": [repr(e)]}
    inv["bmc"] = bmc
    log(f"bmc: {json.dumps({k: v for k, v in (bmc or {}).items() if k in ('source', 'address', 'user', 'configured', 'verified_oob', 'errors')})}")

    sysinfo = inv.get("system") or {}
    inv["identity"] = {"uuid": sysinfo.get("system-uuid"), "serial": sysinfo.get("system-serial-number"),
                       "macs": macs, "boot_mac": None}
    # the interface that got DHCP is the one live-boot used
    for n in inv.get("network") or []:
        if n.get("ipv4"):
            inv["identity"]["boot_mac"] = n["mac"]
            break
    inv["duration_s"] = round(time.time() - started, 1)

    resp = post(server, inv)
    if resp:
        log(f"server accepted: machine {resp.get('id')} state={resp.get('state')}")
    else:
        log("could not deliver inventory to server")
        keep = json.loads(json.dumps(inv))
        if (keep.get("bmc") or {}).get("password"):
            keep["bmc"]["password"] = "<not kept>"
        with open(os.open("/root/inventory.json", os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600), "w") as f:
            json.dump(keep, f, indent=1)
    if opts.get("nopoweroff") or not resp:
        log("staying up (nopoweroff or delivery failure)")
        return 0 if resp else 1
    log("powering off")
    time.sleep(3)
    run(["systemctl", "poweroff"])
    return 0


if __name__ == "__main__":
    sys.exit(main())