infix/board/common/rootfs/libexec/infix/probe

#!/usr/bin/env python3

import importlib.machinery
import json
import os
import shutil
import struct
import subprocess
import sys

onieprom = importlib.machinery.SourceFileLoader("onieprom","/bin/onieprom").load_module()
SYSTEM_JSON   = "/run/system.json"
KKIT_IANA_PEM = 61046

class DTSystem:
    BASE = "/sys/firmware/devicetree/base"
    INFIX = BASE + "/chosen/infix"

    def __init__(self):
        self.vpdseq = 0

        dt = {}
        for root, _, files in os.walk(DTSystem.BASE):
            if "phandle" not in files:
                continue

            phandle = os.path.join(root, "phandle")
            if not os.path.exists(phandle):
                continue

            ph, = struct.unpack(">L", open(phandle, "rb").read())
            dt[ph] = root

        sys = {}
        for root, dirs, _ in os.walk("/sys/devices"):
            if "of_node" not in dirs:
                continue

            phandle = os.path.join(root, "of_node", "phandle")
            if not os.path.exists(phandle):
                continue

            ph, = struct.unpack(">L", open(phandle, "rb").read())
            if ph not in sys:
                sys[ph] = []
            sys[ph].append(root)

        phs = set(list(dt.keys()) + list(sys.keys()))

        self.devs = { ph: [Device(ph, dt.get(ph), s if s is not None else "") for s in (sys.get(ph) or []) if ph is not None]  for ph in phs }
        self.base = Device(0, None, DTSystem.BASE)
        self.infix = Device(0, None, DTSystem.INFIX)

    def __get_phandle_array(self, name):
        path = os.path.join(DTSystem.INFIX, name)
        if not os.path.exists(path):
            return ()

        data = open(path, "rb").read()
        elems = len(data) // struct.calcsize(">L")
        return struct.unpack(">" + elems * "L", data)

    def devices_from_ph(self, ph):
        return self.devs.get(ph)

    def into_vpd(self, dev):
        def parse():
            if not dev.available():
                return {}

            try:
                data = onieprom.from_tlv(open(dev.attrpath("nvmem"), "rb", 0))
            except:
                data = {}

            return data

        self.vpdseq += 1
        return {
            "board": dev.dtstr("infix,board", f"UNKNOWN{self.vpdseq}"),
            "available": dev.available(),
            "trusted": dev.hasdtattr("infix,trusted"),
            "data": parse(),
        }

    def infix_usb_devices(self, out):
        names=self.infix.str_array("usb-port-names", ())
        phs=self.__get_phandle_array("usb-ports")
        data=dict(zip(names,phs))
        if data != {}:
            out["usb-ports"] = []
        for name,ph in data.items():
            [out["usb-ports"].extend([{"name": name, "path": dev.attrpath("authorized")}, {"name": name, "path": dev.attrpath("authorized_default")}]) for dev in self.devices_from_ph(ph)]


    def infix_devices(self, kind):
        phs = self.__get_phandle_array(kind)
        return [[d for d in self.devices_from_ph(ph)] for ph in phs]

    def infix_vpds(self):
        flat_devices = [device for sublist in self.infix_devices("vpds") for device in sublist]
        return [self.into_vpd(device) for device in flat_devices]

class QEMUSystem:
    BASE = "/sys/firmware/qemu_fw_cfg"
    REV = BASE + "/rev"
    VPD = BASE + "/by_name/opt/vpd/raw"

    def product_vpd(self):
        data = {}
        if os.path.exists(QEMUSystem.VPD):
            try:
                data = onieprom.from_tlv(open(QEMUSystem.VPD, "rb", 0))
            except:
                pass

        return {
            "board": "product",
            "available": os.path.exists(QEMUSystem.VPD),
            "trusted": True,
            "data": data,
        }

    def vpds(self):
        return [self.product_vpd()]

    def usb_ports(self):
        ports = [
            {
            "name": "USB",
            "path": "/sys/bus/usb/devices/usb1/authorized"

        },
        {
            "name": "USB",
            "path": "/sys/bus/usb/devices/usb1/authorized_default"

        },
        {
            "name": "USB2",
            "path": "/sys/bus/usb/devices/usb2/authorized"
        },
        {
            "name": "USB2",
            "path": "/sys/bus/usb/devices/usb2/authorized_default"

        }]
        return ports
class Device:
    def __init__(self, ph, dtpath, syspath):
        self.ph, self.dtpath, self.syspath = ph, dtpath, syspath

    def available(self):
        return self.syspath != None

    def __getitem__(self, attr):
        return self.attr(attr).decode("utf-8").strip("\0")

    def __setitem__(self, attr, value):
        return self.attr(attr, val=value.encode("utf-8"))


    def attrpath(self, attr):
        return os.path.join(self.syspath, attr)

    def hasattr(self, attr):
        return os.path.exists(self.attrpath(attr))

    def attr(self, attr, default=None, val=None):
        if not self.hasattr(attr):
            return default if val == None else False

        if val:
            open(self.attrpath(attr), "wb").write(value)
            return True

        return open(self.attrpath(attr), "rb").read()

    def str(self, attr, default=None):
        val = self.attr(attr)
        return val.decode("utf-8").strip("\0") if val else default

    def str_array(self, attr, default=None):
        val = self.attr(attr)
        return val.decode("utf-8").strip("\0").split("\0") if val else default

    def dtattrpath(self, attr):
        return os.path.join(self.dtpath, attr)

    def hasdtattr(self, attr):
        return os.path.exists(self.dtattrpath(attr))

    def dtattr(self, attr, default=None):
        if not self.hasdtattr(attr):
            return default

        return open(self.dtattrpath(attr), "rb").read()

    def dtstr(self, attr, default=None):
        val = self.dtattr(attr)
        return val.decode("utf-8").strip("\0") if val else default


def vpd_get_json_ve(vpd, pem):
    ves = vpd["data"].get("vendor-extension")
    if not ves:
        return {}

    out = {}
    for ve in filter(lambda ve: ve[0] == pem, ves):
        out.update(json.loads(ve[1]))

    return out

def vpd_get_pwhash(vpd):
    if not vpd.get("trusted"):
        return None

    kkit = vpd_get_json_ve(vpd, KKIT_IANA_PEM)
    return kkit.get("pwhash")

def vpd_inject(out, vpds):
    out["vpd"] = { vpd["board"]: vpd for vpd in vpds }

    product = out["vpd"].get("product", {}).get("data", {})
    hoistattrs = ("vendor", "product-name", "part-number", "serial-number", "mac-address")
    for attr in hoistattrs:
        if attr in product:
            out[attr] = product[attr]

    for vpd in vpds:
        pwhash = vpd_get_pwhash(vpd)
        if pwhash:
            out["factory-password-hash"] = pwhash
            break

def qemu_base_mac():
    """Find MAC address of first non-loopback interface, subtract with 1"""
    base_path = '/sys/class/net'
    interfaces = []

    for iface in os.listdir(base_path):
        if iface == 'lo':
            continue
        try:
            # pylint: disable=invalid-name
            with open(os.path.join(base_path, iface, 'address'), 'r', encoding='ascii') as f:
                mac = f.read().strip()
            interfaces.append((mac, iface))
        except FileNotFoundError:
            continue

    if interfaces:
        interfaces.sort()
        mac  = interfaces[0][0]
        mac  = int(mac.replace(':', ''), 16)
        mac -= 1
        mac %= 1 << 48
        mac  = ':'.join(f"{(mac >> 8 * i) & 0xff:02x}" for i in range(5, -1, -1))
        return mac

    return None

def probe_qemusystem(out):
    """Probe Qemu based test systems and 'make run'"""
    admin_hash = "$5$mI/zpOAqZYKLC2WU$i7iPzZiIjOjrBF3NyftS9CCq8dfYwHwrmUK097Jca9A"

    qsys = QEMUSystem()
    vpds = qsys.vpds()
    usb_ports = qsys.usb_ports()
    vpd_inject(out, vpds)
    out["usb-ports"] = usb_ports
    for (attr, default) in (
            ("vendor", "QEMU"),
            ("product-name", "VM"),
            ("mac-address", qemu_base_mac()),
    ):
        if not out[attr]:
            out[attr] = default

    if not out["factory-password-hash"] and \
       not out["vpd"]["product"]["available"]:
        # Virtual instance without VPD emulation, fallback to
        # admin/admin
        out["factory-password-hash"] = admin_hash

    # Let others react to the fact that we are running in QEMU
    subprocess.run("initctl -nbq cond set qemu".split(), check=False)
    return 0

def probe_dtsystem(out):
    """Probe DTS based system, expects a VPD in ONIE PROM format."""
    dtsys = DTSystem()
    vpds = dtsys.infix_vpds()
    dtsys.infix_usb_devices(out)
    model = dtsys.base.str("model")
    if model:
        out["product-name"] = model

    staticpw = dtsys.infix.str("factory-password-hash")
    if not out["factory-password-hash"]:
        out["factory-password-hash"] = staticpw

    vpd_inject(out, vpds)
    return 0

def main():
    out = {
        "vendor": None,
        "product-name": None,
        "part-number": None,
        "serial-number": None,
        "mac-address": None,
        "factory-password-hash": None,
        "vpd": {}
    }
    vpds = []

    if os.path.exists(QEMUSystem.REV):
        err = probe_qemusystem(out)
    elif os.path.exists(DTSystem.BASE):
        err = probe_dtsystem(out)
    else:
        return 1

    if err:
        return err

    if not out["factory-password-hash"]:
        sys.stdout.write("\n\n\033[31mCRITICAL BOOTSTRAP ERROR\n" +
                         "NO FACTORY PASSWORD FOUND\033[0m\n\n")
        err = 1

    os.umask(0o337)
    # pylint: disable=invalid-name
    with open(SYSTEM_JSON, "w", encoding="ascii") as f:
        json.dump(out, f)
    shutil.chown(SYSTEM_JSON, user="root", group="wheel")
    return err

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