#!/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 hashlib
import json
import os
import re
import secrets
import socket
import ssl
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 tls_setup(opts, server):
    """rollcall.tls=URL rollcall.ca=SHA256: the CA certificate is fetched and
    accepted only by that fingerprint; from then on everything goes over HTTPS
    with it. When TLS was asked for and does not work, nothing falls back to HTTP."""
    url, want = opts.get("tls"), str(opts.get("ca") or "").replace(":", "").lower()
    if not isinstance(url, str) or not url.startswith("https://") or len(want) != 64:
        return server, None, None
    pem = None
    for i in range(10):
        try:
            with urllib.request.urlopen(f"{server}/ca.pem", timeout=10) as r:
                pem = r.read().decode()
            break
        except (urllib.error.URLError, OSError) as e:
            log(f"ca.pem: {e}, retry {i + 1}/10")
            time.sleep(3)
    if pem is None:
        raise RuntimeError("cannot fetch the CA certificate")
    if hashlib.sha256(ssl.PEM_cert_to_DER_cert(pem)).hexdigest() != want:
        raise RuntimeError("the CA certificate does not match rollcall.ca: not talking to this server")
    path = "/run/rollcall-ca.pem"
    with open(path, "w") as f:
        f.write(pem)
    log(f"talking to {url} over HTTPS (CA {want[:16]}...)")
    return url.rstrip("/"), ssl.create_default_context(cafile=path), path


TLS = None  # ssl context once tls_setup ran


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, context=TLS) 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"), "errors": []}
    if not h.get("password"):
        # the region keeps the credentials and checks the BMC itself (rollcall >= 0.3)
        res["configured"] = True
        return res
    res["password"] = h["password"]
    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, context=TLS) 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, TLS
    server, TLS, _ = tls_setup(opts, server)
    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())