#!/usr/bin/env python3
"""Collect a bounded, read-only Linux PCI evidence snapshot for Header7F.

SPDX-License-Identifier: MIT

By default, read fixed textual sysfs attributes and PCI topology only. An
explicit --include-config request reads at most 4096 bytes from the selected
device's config file. Config reads can affect faulty devices; see --help.
Never read MMIO, process memory, or kernel logs, invoke commands, make network
requests, or change device or kernel settings. No root escalation is attempted.
"""

import argparse
from datetime import datetime, timezone
import json
import os
from pathlib import Path
import re
import sys


SCHEMA = "header7f.linux-evidence.v1"
VERSION = "0.2.0"
LICENSE = "MIT"
MAX_TEXT_BYTES = 4096
MAX_MSI_IRQS = 2048
MAX_TOPOLOGY_DEVICES = 512
MAX_UPSTREAM_PORTS = 16
MAX_ROOT_PORTS = 32
MAX_CONFIG_BYTES = 4096
MAX_OUTPUT_BYTES = 1048576
MAX_WARNINGS = 512
_BDF = re.compile(r"(?:(?P<domain>[0-9a-fA-F]{4,5}):)?(?P<bus>[0-9a-fA-F]{2}):(?P<slot>[0-1][0-9a-fA-F])\.(?P<function>[0-7])", re.ASCII)
_IRQ_NAME = re.compile(r"[0-9]{1,10}", re.ASCII)
_ROOT_BUS = re.compile(r"pci[0-9a-fA-F]{4,5}:[0-9a-fA-F]{2}", re.ASCII)
LINK_FIELDS = ("current_link_speed", "current_link_width", "max_link_speed", "max_link_width")
AER_FIELDS = (
    "aer_dev_correctable", "aer_dev_fatal", "aer_dev_nonfatal",
    "aer_rootport_total_err_cor", "aer_rootport_total_err_fatal", "aer_rootport_total_err_nonfatal",
)


def normalize_bdf(value):
    """Accept a complete BDF or a short BDF in domain 0000, never a path."""
    if not isinstance(value, str):
        raise ValueError("Use a PCI address such as 0000:01:00.0 or 01:00.0.")
    match = _BDF.fullmatch(value)
    if match is None:
        raise ValueError("Use a PCI address such as 0000:01:00.0 or 01:00.0.")
    return "{domain}:{bus}:{slot}.{function}".format(
        domain=match.group("domain") or "0000",
        bus=match.group("bus"),
        slot=match.group("slot"),
        function=match.group("function"),
    ).lower()


def _read_text(path, label, warnings, max_bytes=MAX_TEXT_BYTES):
    try:
        with path.open("rb") as source:
            value = source.read(max_bytes + 1)
    except OSError:
        warnings.append("{} is unavailable or unreadable.".format(label))
        return None
    if len(value) > max_bytes:
        warnings.append("{} exceeds the collection limit and was omitted.".format(label))
        return None
    try:
        text = value.decode("utf-8").strip()
    except UnicodeError:
        warnings.append("{} is not valid text and was omitted.".format(label))
        return None
    if not text or "\x00" in text:
        warnings.append("{} is empty or not valid text.".format(label))
        return None
    return text


def _read_link_name(path, label, warnings):
    try:
        if path.is_symlink():
            return path.resolve(strict=True).name
    except (OSError, RuntimeError):
        pass
    warnings.append("{} is unavailable or has no readable link.".format(label))
    return None


def _read_msi_irqs(device_path, warnings):
    entries = []
    try:
        with os.scandir(device_path / "msi_irqs") as directory:
            for index, entry in enumerate(directory):
                if len(entries) >= MAX_MSI_IRQS or index >= MAX_MSI_IRQS * 2:
                    warnings.append("MSI IRQ collection reached its limit; the vector list may be incomplete.")
                    break
                if _IRQ_NAME.fullmatch(entry.name) is None:
                    warnings.append("An unexpected MSI IRQ entry was omitted.")
                    continue
                entries.append(entry.name)
    except OSError:
        warnings.append("msi_irqs is unavailable or unreadable; an empty list does not prove MSI is disabled.")
        return []
    result = []
    for irq in sorted(entries, key=lambda value: (int(value), value)):
        mode = _read_text(device_path / "msi_irqs" / irq, "MSI IRQ {} mode".format(irq), warnings, 64)
        if mode is not None and mode not in ("msi", "msix"):
            warnings.append("MSI IRQ {} has an unrecognized mode.".format(irq))
            mode = None
        result.append({"irq": irq, "mode": mode})
    return result


def _read_config(device_path, requested, present, warnings):
    """Read only the selected target, only on explicit opt-in, without padding."""
    result = {
        "requested": requested, "status": "not-requested", "requestedBytes": MAX_CONFIG_BYTES,
        "bytesRead": 0, "hex": None, "error": None,
    }
    if not requested:
        return result
    if not present:
        result.update(status="target-absent", error="The selected target is not available in sysfs.")
        return result
    try:
        with (device_path / "config").open("rb") as source:
            value = source.read(MAX_CONFIG_BYTES)
    except OSError as error:
        # Do not include arbitrary OS error text or full host paths in an export.
        result.update(status="unreadable", error="Permission denied." if isinstance(error, PermissionError) else "Config space is unavailable or unreadable.")
        warnings.append("Selected target config space could not be read; no bytes were inferred.")
        return result
    result.update(bytesRead=len(value), hex=value.hex())
    if not value:
        result.update(status="empty", error="The config read returned no bytes.")
        warnings.append("Selected target config space returned no bytes.")
    elif len(value) < MAX_CONFIG_BYTES:
        result.update(status="partial", error="Only {} of {} requested bytes were readable; remaining bytes are unknown.".format(len(value), MAX_CONFIG_BYTES))
        warnings.append("Selected target config space is partial; unread bytes were not padded or inferred.")
    else:
        result["status"] = "complete"
    return result


def _capture_metadata(address, device_path, warnings, include_config=False, link_aer_only=False):
    present = device_path.is_dir()
    device = {"bdf": address, "present": present}
    if not present:
        for name in ("vendor", "device", "class", "driver", "resource", *LINK_FIELDS, "iommu_group", "irq", *AER_FIELDS):
            device[name] = None
        device["msi_irqs"] = []
    else:
        for name in ("vendor", "device", "class"):
            device[name] = _read_text(device_path / name, "{} {}".format(address, name), warnings, 64)
        device["driver"] = _read_link_name(device_path / "driver", "{} driver".format(address), warnings)
        device["resource"] = None if link_aer_only else _read_text(device_path / "resource", "{} resource".format(address), warnings)
        for name in LINK_FIELDS:
            device[name] = _read_text(device_path / name, "{} {}".format(address, name), warnings, 128)
        device["iommu_group"] = None if link_aer_only else _read_link_name(device_path / "iommu_group", "{} iommu_group".format(address), warnings)
        device["irq"] = None if link_aer_only else _read_text(device_path / "irq", "{} irq".format(address), warnings, 64)
        device["msi_irqs"] = [] if link_aer_only else _read_msi_irqs(device_path, warnings)
        for name in AER_FIELDS:
            device[name] = _read_text(device_path / name, "{} {}".format(address, name), warnings)
    device["config"] = _read_config(device_path, include_config, present, warnings)
    if link_aer_only:
        device["detailScope"] = "link-aer-only"
    return device


def _pci_parents(device_path, warnings):
    """Use the real sysfs hierarchy; bus numbers alone do not locate a parent."""
    try:
        resolved = device_path.resolve(strict=True)
    except (OSError, RuntimeError):
        warnings.append("PCI hierarchy for {} is unavailable.".format(device_path.name))
        return []
    parents = []
    for parent in resolved.parents:
        if _BDF.fullmatch(parent.name) and parent.name.count(":") == 2:
            parents.append(parent)
            if len(parents) >= MAX_UPSTREAM_PORTS:
                warnings.append("PCI upstream chain reached its limit; earlier ancestors may be omitted.")
                break
        if _ROOT_BUS.fullmatch(parent.name):
            break
    return parents


def _collect_topology(pci_devices, target_address, target_present, explicit_upstream, warnings):
    topology = {
        "devices": [], "upstream": [], "rootPorts": [], "targetParentBdf": None,
        "targetParentSource": None, "complete": True,
    }
    names = []
    try:
        with os.scandir(pci_devices) as entries:
            for index, entry in enumerate(entries):
                if index >= MAX_TOPOLOGY_DEVICES * 2 or len(names) >= MAX_TOPOLOGY_DEVICES:
                    topology["complete"] = False
                    warnings.append("PCI topology reached its collection limit; some devices may be omitted.")
                    break
                if entry.name.count(":") == 2 and _BDF.fullmatch(entry.name):
                    names.append(entry.name)
    except OSError:
        topology["complete"] = False
        warnings.append("PCI device enumeration is unavailable; an empty topology does not prove no devices exist.")
    root_candidates = []
    for name in sorted(names):
        path = pci_devices / name
        if not path.is_dir():
            topology["complete"] = False
            warnings.append("Topology entry {} became unavailable during collection.".format(name))
            continue
        parents = _pci_parents(path, warnings)
        parent_bdf = parents[0].name.lower() if parents else None
        node = {"bdf": name.lower(), "parentBdf": parent_bdf}
        for field in ("class", "vendor", "device"):
            node[field] = _read_text(path / field, "{} topology {}".format(name, field), warnings, 64)
        topology["devices"].append(node)
        # A top-level PCI bridge is only a root-port candidate without config capability decoding.
        if parent_bdf is None and node["class"] and node["class"].lower().startswith("0x0604"):
            root_candidates.append(name)
    if len(root_candidates) > MAX_ROOT_PORTS:
        topology["complete"] = False
        warnings.append("Root bridge candidates reached their collection limit.")
    metadata_cache = {}

    def metadata(address):
        if address not in metadata_cache:
            metadata_cache[address] = _capture_metadata(address, pci_devices / address, warnings, link_aer_only=True)
        return metadata_cache[address]

    topology["rootPorts"] = [metadata(name) for name in root_candidates[:MAX_ROOT_PORTS]]
    if target_present:
        parents = _pci_parents(pci_devices / target_address, warnings)
        if len(parents) >= MAX_UPSTREAM_PORTS:
            topology["complete"] = False
        if parents:
            topology["targetParentBdf"] = parents[0].name.lower()
            topology["targetParentSource"] = "sysfs"
            topology["upstream"] = [metadata(parent.name.lower()) for parent in parents]
        if explicit_upstream and explicit_upstream != topology["targetParentBdf"]:
            warnings.append("The supplied upstream address does not match the target's sysfs parent; the sysfs hierarchy was retained.")
    elif explicit_upstream:
        topology["targetParentBdf"] = explicit_upstream
        topology["targetParentSource"] = "user-specified"
        topology["upstream"] = [metadata(explicit_upstream)]
        if topology["upstream"][0]["present"]:
            parents = _pci_parents(pci_devices / explicit_upstream, warnings)
            if len(parents) >= MAX_UPSTREAM_PORTS:
                topology["complete"] = False
            topology["upstream"].extend(metadata(parent.name.lower()) for parent in parents[:MAX_UPSTREAM_PORTS - 1])
        warnings.append("The supplied upstream device is context only: an absent target's physical connection cannot be verified from its address.")
    return topology


def _bound_capture(evidence):
    """Keep the actual indented ASCII JSON within the browser import limit."""
    if len(evidence["warnings"]) > MAX_WARNINGS:
        evidence["warnings"] = evidence["warnings"][:MAX_WARNINGS - 1] + ["Additional collection limitations were omitted after the warning limit."]

    def size():
        return len(json.dumps(evidence, indent=2, ensure_ascii=True, allow_nan=False).encode("utf-8")) + 1

    if size() > MAX_OUTPUT_BYTES:
        evidence["topology"]["complete"] = False
        message = "Additional topology was omitted to keep the capture within the 1 MiB browser import limit."
        if len(evidence["warnings"]) >= MAX_WARNINGS:
            evidence["warnings"] = evidence["warnings"][:MAX_WARNINGS - 1]
            message = "Further collection limitations and additional topology were omitted to keep the capture within the 1 MiB browser import limit."
        evidence["warnings"].append(message)
        # Retain the selected target and nearest upstream information ahead of
        # unrelated root candidates or a broad device list.
        for name in ("rootPorts", "devices", "upstream"):
            while evidence["topology"][name] and size() > MAX_OUTPUT_BYTES:
                evidence["topology"][name].pop()
    return evidence


def capture_device(bdf, sysfs_root="/sys", include_config=False, upstream_bdf=None):
    """Return evidence for one device; sysfs_root exists for fixture tests only.

    The command line intentionally has no alternate-root option.
    """
    address = normalize_bdf(bdf)
    upstream = normalize_bdf(upstream_bdf) if upstream_bdf is not None else None
    if upstream == address:
        raise ValueError("The upstream address must differ from the selected target.")
    pci_devices = Path(sysfs_root) / "bus" / "pci" / "devices"
    device_path = pci_devices / address
    warnings = []
    device = _capture_metadata(address, device_path, warnings, include_config)
    if not device["present"]:
        warnings.append("PCI target {} is absent or unavailable in sysfs at collection time; this does not identify why it is missing.".format(address))
    topology = _collect_topology(pci_devices, address, device["present"], upstream, warnings)
    return _bound_capture({
        "schema": SCHEMA,
        "collector": {"name": "header7f-linux-capture", "version": VERSION, "license": LICENSE, "configReadRequested": include_config},
        "collectedAt": datetime.now(timezone.utc).isoformat(timespec="milliseconds").replace("+00:00", "Z"),
        "device": device,
        "topology": topology,
        "kernelLog": {"provided": False, "text": ""},
        "warnings": warnings,
    })


def main(argv=None):
    parser = argparse.ArgumentParser(
        description="Collect read-only Linux PCI metadata as a Header7F JSON capture.",
        epilog=(
            "Example: python3 header7f-linux-capture.py 0000:01:00.0 > capture.json\n"
            "Short addresses such as 01:00.0 use domain 0000; 4- or 5-digit domains are accepted.\n"
            "Default collection needs no sudo: bounded textual sysfs metadata, PCI topology, max/current link values, and AER counters.\n"
            "An absent target still produces a capture of the visible PCI topology and root bridge candidates.\n"
            "--upstream selects a known upstream device when the target is absent; its connection is not inferred or verified.\n"
            "--include-config explicitly reads up to 4096 config-space bytes from the selected target only. "
            "Some devices can crash on config reads; use only when appropriate for your hardware. "
            "Permission failures and partial reads remain explicit; no root escalation is attempted.\n"
            "No hardware writes, MMIO reads, kernel logs, external commands, network calls, or kernel changes.\n"
            "Address allocation does not prove that DMA works or identify DMA buffer mappings. "
            "AER counters are cumulative reports at their device, not proof of the cause or timing of this failure. "
            "The UTC timestamp marks collection; sequential fields are not an atomic trace. "
            "Review the JSON before importing or sharing it."
        ),
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("--version", action="version", version="%(prog)s {} ({}; {})".format(VERSION, LICENSE, SCHEMA))
    parser.add_argument("--include-config", action="store_true", help="explicitly read up to 4096 target config bytes; some devices can crash on config reads")
    parser.add_argument("--upstream", metavar="BDF", help="known upstream BDF for a missing target; retained as user-supplied context")
    parser.add_argument("bdf", help="PCI device address, for example 0000:01:00.0")
    args = parser.parse_args(argv)
    if not sys.platform.startswith("linux"):
        parser.exit(1, "header7f: this collector requires Linux.\n")
    try:
        evidence = capture_device(args.bdf, include_config=args.include_config, upstream_bdf=args.upstream)
    except (ValueError, OSError) as error:
        parser.exit(1, "header7f: {}\n".format(error))
    json.dump(evidence, sys.stdout, indent=2, ensure_ascii=True, allow_nan=False)
    sys.stdout.write("\n")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())