PCILeechFWGenerator 0.5.7__py3-none-any.whl

cli/__init__.py

@@ -0,0 +1,17 @@
+"""CLI components for PCILeech FW Generator."""
+
+from .cli import get_parser, main
+from .config import BuildConfig
+from .container import require_podman, run_build
+from .flash import flash_firmware
+from .vfio_handler import VFIOBinder
+
+__all__ = [
+    "get_parser",
+    "main",
+    "BuildConfig",
+    "run_build",
+    "require_podman",
+    "flash_firmware",
+    "VFIOBinder",
+]

cli/build_wrapper.py

@@ -0,0 +1,41 @@
+#!/usr/bin/env python3
+"""
+Wrapper script to properly invoke the build.py module with correct Python path setup.
+This ensures that relative imports work correctly in the container environment.
+"""
+
+import os
+import sys
+from pathlib import Path
+
+# Determine if we're in container or local environment
+if Path("/app").exists():
+    # Container environment
+    app_dir = Path("/app")
+    src_dir = app_dir / "src"
+else:
+    # Local environment - we're in src/cli, so go up two levels to get to project root
+    app_dir = Path(__file__).parent.parent.parent.absolute()
+    src_dir = app_dir / "src"
+
+# Ensure both directories are in Python path
+if str(app_dir) not in sys.path:
+    sys.path.insert(0, str(app_dir))
+if str(src_dir) not in sys.path:
+    sys.path.insert(0, str(src_dir))
+
+# Change to the src directory so relative imports work
+if src_dir.exists():
+    os.chdir(str(src_dir))
+else:
+    print(f"Error: Source directory {src_dir} does not exist")
+    sys.exit(1)
+
+# Now import and run the build module
+if __name__ == "__main__":
+    # Import the build module
+    import build
+
+    # Run the main function with the original arguments
+    sys.argv[0] = "build.py"  # Fix the script name for argument parsing
+    build.main()

cli/cli.py

@@ -0,0 +1,222 @@
+#!/usr/bin/env python3
+"""cli – one front‑door for the whole tool‑chain.
+
+Usage examples
+~~~~~~~~~~~~~~
+    # guided build flow (device & board pickers)
+    ./cli build
+
+    # scripted build for CI (non‑interactive)
+    ./cli build --bdf 0000:01:00.0 --board pcileech_75t484_x1 \
+                --device-type network --advanced-sv
+
+    # flash an already‑generated bitstream
+    ./cli flash output/firmware.bin --board pcileech_75t484_x1
+"""
+from __future__ import annotations
+
+import argparse
+import logging
+import os
+import re
+import subprocess
+import sys
+from pathlib import Path
+from typing import Dict, List, Optional
+
+from log_config import get_logger, setup_logging
+from shell import Shell
+
+from .container import BuildConfig, run_build  # new unified runner
+
+logger = get_logger(__name__)
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Helpers – PCIe enumeration & interactive pickers
+# ──────────────────────────────────────────────────────────────────────────────
+PCI_RE = re.compile(
+    r"(?P<bdf>[0-9a-fA-F:.]+) .*?\["
+    r"(?P<class>[0-9a-fA-F]{4})\]: .*?\["
+    r"(?P<ven>[0-9a-fA-F]{4}):(?P<dev>[0-9a-fA-F]{4})\]"
+)
+
+
+def list_pci_devices() -> List[Dict[str, str]]:
+    out = Shell().run("lspci -Dnn")
+    devs: list[dict[str, str]] = []
+    for line in out.splitlines():
+        m = PCI_RE.match(line)
+        if m:
+            d = m.groupdict()
+            d["pretty"] = line
+            devs.append(d)
+    return devs
+
+
+def pick(lst: list[str], prompt: str) -> str:
+    for i, item in enumerate(lst):
+        print(f" [{i}] {item}")
+    while True:
+        sel = input(prompt).strip()
+        if not sel and lst:
+            return lst[0]
+        try:
+            return lst[int(sel)]
+        except Exception:
+            print("  Invalid selection – try again.")
+
+
+def choose_device() -> Dict[str, str]:
+    devs = list_pci_devices()
+    if not devs:
+        raise RuntimeError("No PCIe devices found – are you root?")
+    for i, dev in enumerate(devs):
+        print(f" [{i}] {dev['pretty']}")
+    return devs[int(input("Select donor device #: "))]
+
+
+SUPPORTED_BOARDS = [
+    "pcileech_75t484_x1",
+    "pcileech_35t484_x1",
+    "pcileech_35t325_x4",
+    "pcileech_35t325_x1",
+    "pcileech_100t484_x1",
+    "pcileech_enigma_x1",
+    "pcileech_squirrel",
+    "pcileech_pciescreamer_xc7a35",
+]
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# CLI setup
+# ──────────────────────────────────────────────────────────────────────────────
+
+
+def build_sub(parser: argparse._SubParsersAction):
+    p = parser.add_parser("build", help="Build firmware (guided or scripted)")
+    p.add_argument("--bdf", help="PCI BDF (skip for interactive picker)")
+    p.add_argument("--board", choices=SUPPORTED_BOARDS, help="FPGA board")
+    p.add_argument(
+        "--device-type",
+        default="generic",
+        choices=["generic", "network", "storage", "graphics", "audio"],
+        help="Type of device being cloned",
+    )
+    p.add_argument(
+        "--advanced-sv", action="store_true", help="Enable advanced SV features"
+    )
+    p.add_argument("--enable-variance", action="store_true", help="Enable variance")
+    p.add_argument(
+        "--auto-fix", action="store_true", help="Let VFIOBinder auto-remediate issues"
+    )
+
+    # Add fallback control group
+    fallback_group = p.add_argument_group("Fallback Control")
+    fallback_group.add_argument(
+        "--fallback-mode",
+        choices=["none", "prompt", "auto"],
+        default="none",
+        help="Control fallback behavior (none=fail-fast, prompt=ask, auto=allow)",
+    )
+    fallback_group.add_argument(
+        "--allow-fallbacks", type=str, help="Comma-separated list of allowed fallbacks"
+    )
+    fallback_group.add_argument(
+        "--deny-fallbacks", type=str, help="Comma-separated list of denied fallbacks"
+    )
+    fallback_group.add_argument(
+        "--legacy-compatibility",
+        action="store_true",
+        help="Enable legacy compatibility mode (temporarily restores old fallback behavior)",
+    )
+
+
+def flash_sub(parser: argparse._SubParsersAction):
+    p = parser.add_parser("flash", help="Flash a firmware binary via usbloader")
+    p.add_argument("firmware", help="Path to .bin")
+    p.add_argument(
+        "--board", required=True, choices=SUPPORTED_BOARDS, help="FPGA board"
+    )
+
+
+def get_parser() -> argparse.ArgumentParser:
+    ap = argparse.ArgumentParser("cli", description=__doc__)
+    sub = ap.add_subparsers(
+        dest="cmd",
+        required=True,
+        help="Command to run (build/flash)",
+    )
+    build_sub(sub)
+    flash_sub(sub)
+    return ap
+
+
+def flash_bin(path: Path):
+    from .flash import flash_firmware
+
+    flash_firmware(path)
+
+    logger.info("Firmware flashed successfully ✓")
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Entrypoint
+# ──────────────────────────────────────────────────────────────────────────────
+
+
+def main(argv: Optional[List[str]] = None):
+    # Setup proper logging with color support
+    setup_logging(level=logging.INFO)
+
+    args = get_parser().parse_args(argv)
+
+    if args.cmd == "build":
+        bdf = args.bdf or choose_device()["bdf"]
+        board = args.board or pick(SUPPORTED_BOARDS, "Board #: ")
+        # Process fallback lists
+        allowed_fallbacks = []
+        if hasattr(args, "allow_fallbacks") and args.allow_fallbacks:
+            allowed_fallbacks = [f.strip() for f in args.allow_fallbacks.split(",")]
+
+        denied_fallbacks = []
+        if hasattr(args, "deny_fallbacks") and args.deny_fallbacks:
+            denied_fallbacks = [f.strip() for f in args.deny_fallbacks.split(",")]
+
+        # Determine fallback mode based on legacy compatibility flag
+        fallback_mode = getattr(args, "fallback_mode", "none")
+        if (
+            hasattr(args, "legacy_compatibility")
+            and args.legacy_compatibility
+            and fallback_mode == "none"
+        ):
+            logger.warning(
+                "Legacy compatibility mode enabled - using 'auto' fallback mode"
+            )
+            fallback_mode = "auto"
+            if not allowed_fallbacks:
+                allowed_fallbacks = [
+                    "config-space",
+                    "msix",
+                    "behavior-profiling",
+                    "build-integration",
+                ]
+
+        cfg = BuildConfig(
+            bdf=bdf,
+            board=board,
+            device_type=args.device_type,
+            advanced_sv=args.advanced_sv,
+            enable_variance=args.enable_variance,
+            auto_fix=args.auto_fix,
+            fallback_mode=fallback_mode,
+            allowed_fallbacks=allowed_fallbacks,
+            denied_fallbacks=denied_fallbacks,
+        )
+        run_build(cfg)
+
+    elif args.cmd == "flash":
+        flash_bin(Path(args.firmware))
+
+
+if __name__ == "__main__":
+    main()

cli/config.py

@@ -0,0 +1,66 @@
+"""Configuration dataclass for PCILeech firmware generation."""
+
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class BuildConfig:
+    """Strongly-typed configuration for firmware build process."""
+
+    # Device configuration
+    bdf: str
+    vendor: str
+    device: str
+    board: str
+    device_type: str
+
+    # Advanced features
+    advanced_sv: bool = True
+    enable_variance: bool = True
+    donor_dump: bool = True
+    auto_install_headers: bool = True
+    strict_vfio: bool = True  # Fail hard if VFIO is not available
+
+    # Feature toggles
+    disable_power_management: bool = False
+    disable_error_handling: bool = False
+    disable_performance_counters: bool = False
+
+    flash: bool = True
+
+    # Timing configuration
+    behavior_profile_duration: int = 45
+
+    # Mode configuration
+    tui: bool = False
+    interactive: bool = False
+
+    # Runtime state
+    original_driver: Optional[str] = None
+    iommu_group: Optional[str] = None
+    vfio_device: Optional[str] = None
+
+    def __post_init__(self):
+        """Validate configuration after initialization."""
+
+        # Validate BDF format
+        import re
+
+        bdf_pattern = re.compile(
+            r"^[0-9a-fA-F]{4}:[0-9a-fA-F]{2}:[0-9a-fA-F]{2}\.[0-7]$"
+        )
+        if not bdf_pattern.match(self.bdf):
+            raise ValueError(
+                f"Invalid BDF format: {self.bdf}. Expected format: DDDD:BB:DD.F"
+            )
+
+        # Validate vendor and device IDs
+        if not re.match(r"^[0-9a-fA-F]{4}$", self.vendor):
+            raise ValueError(
+                f"Invalid vendor ID format: {self.vendor}. Expected 4-digit hex."
+            )
+        if not re.match(r"^[0-9a-fA-F]{4}$", self.device):
+            raise ValueError(
+                f"Invalid device ID format: {self.device}. Expected 4-digit hex."
+            )

cli/container.py

@@ -0,0 +1,322 @@
+#!/usr/bin/env python3
+"""container_build – unified VFIO‑aware Podman build runner"""
+from __future__ import annotations
+
+import shutil
+import subprocess
+import sys
+import textwrap
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import List, Optional
+
+from log_config import get_logger
+from shell import Shell
+
+from .vfio_handler import VFIOBinder  # auto‑fix & diagnostics baked in
+
+# Import safe logging functions
+try:
+    from ..string_utils import (log_debug_safe, log_error_safe, log_info_safe,
+                                log_warning_safe)
+except ImportError:
+    # Fallback implementations
+    def log_info_safe(logger, template, **kwargs):
+        logger.info(template.format(**kwargs))
+
+    def log_error_safe(logger, template, **kwargs):
+        logger.error(template.format(**kwargs))
+
+    def log_warning_safe(logger, template, **kwargs):
+        logger.warning(template.format(**kwargs))
+
+    def log_debug_safe(logger, template, **kwargs):
+        logger.debug(template.format(**kwargs))
+
+
+logger = get_logger(__name__)
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Exceptions
+# ──────────────────────────────────────────────────────────────────────────────
+class ContainerError(RuntimeError):
+    pass
+
+
+class VFIOError(RuntimeError):
+    pass
+
+
+class EnvError(RuntimeError):
+    pass
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Build configuration (thin wrapper over original)
+# ──────────────────────────────────────────────────────────────────────────────
+@dataclass
+class BuildConfig:
+    bdf: str
+    board: str
+    # feature toggles
+    advanced_sv: bool = False
+    device_type: str = "generic"
+    enable_variance: bool = False
+    disable_power_management: bool = False
+    disable_error_handling: bool = False
+    disable_performance_counters: bool = False
+    behavior_profile_duration: int = 30
+    # runtime toggles
+    auto_fix: bool = True  # hand to VFIOBinder
+    container_tag: str = "latest"
+    container_image: str = "pcileech-fw-generator"
+    # fallback control options
+    fallback_mode: str = "none"  # "none", "prompt", or "auto"
+    allowed_fallbacks: List[str] = field(default_factory=list)
+    denied_fallbacks: List[str] = field(default_factory=list)
+
+    def cmd_args(self) -> List[str]:
+        """Translate config to build.py flags"""
+        args = [f"--bdf {self.bdf}", f"--board {self.board}"]
+        if self.advanced_sv:
+            args.append("--advanced-sv")
+        if self.device_type != "generic":
+            args.append(f"--device-type {self.device_type}")
+        if self.enable_variance:
+            args.append("--enable-variance")
+        if self.disable_power_management:
+            args.append("--disable-power-management")
+        if self.disable_error_handling:
+            args.append("--disable-error-handling")
+        if self.disable_performance_counters:
+            args.append("--disable-performance-counters")
+        if self.behavior_profile_duration != 30:
+            args.append(f"--behavior-profile-duration {self.behavior_profile_duration}")
+
+        # Add fallback control arguments
+        if self.fallback_mode != "none":
+            args.append(f"--fallback-mode {self.fallback_mode}")
+        if self.allowed_fallbacks:
+            args.append(f"--allow-fallbacks {','.join(self.allowed_fallbacks)}")
+        if self.denied_fallbacks:
+            args.append(f"--deny-fallbacks {','.join(self.denied_fallbacks)}")
+
+        return args
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Helpers
+# ──────────────────────────────────────────────────────────────────────────────
+
+
+def require_podman() -> None:
+    if shutil.which("podman") is None:
+        raise EnvError("Podman not found – install it or adjust PATH")
+
+
+def image_exists(name: str) -> bool:
+    shell = Shell()
+    out = shell.run("podman images --format '{{.Repository}}:{{.Tag}}'", timeout=5)
+    return any(line.startswith(name) for line in out.splitlines())
+
+
+def build_image(name: str, tag: str) -> None:
+    log_info_safe(logger, "Building container image {name}:{tag}", name=name, tag=tag)
+    cmd = f"podman build -t {name}:{tag} -f Containerfile ."
+    subprocess.run(cmd, shell=True, check=True)
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# Public façade
+# ──────────────────────────────────────────────────────────────────────────────
+
+
+def run_build(cfg: BuildConfig) -> None:
+    """High‑level orchestration: VFIO bind → container run → cleanup"""
+    require_podman()
+    if not image_exists(f"{cfg.container_image}:{cfg.container_tag}"):
+        build_image(cfg.container_image, cfg.container_tag)
+
+    # Ensure host output dir exists and is absolute
+    output_dir = (Path.cwd() / "output").resolve()
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    try:
+        # Bind without keeping the FD (call the context manager only long
+        # enough to flip the drivers)
+        binder = VFIOBinder(cfg.bdf, attach=False)
+        with binder:
+            # enter/exit immediately → binds device
+            pass
+
+        # Get the group device path as a string (safe, just a string)
+        from .vfio_handler import _get_iommu_group
+
+        group_id = _get_iommu_group(cfg.bdf)
+        group_dev = f"/dev/vfio/{group_id}"
+
+        log_info_safe(
+            logger,
+            "Launching build container – board={board}, tag={tag}",
+            board=cfg.board,
+            tag=cfg.container_tag,
+            prefix="CONT",
+        )
+
+        cmd_args = " ".join(cfg.cmd_args())
+        podman_cmd = textwrap.dedent(
+            f"""
+            podman run --rm --privileged \
+              --device={group_dev} \
+              --device=/dev/vfio/vfio \
+              --entrypoint python3 \
+              --user root \
+              -v {output_dir}:/app/output \
+              -v /lib/modules/$(uname -r)/build:/kernel-headers:ro \
+              {cfg.container_image}:{cfg.container_tag} \
+              src/build.py {cmd_args}
+            """
+        ).strip()
+
+        log_debug_safe(
+            logger, "Container command: {cmd}", cmd=podman_cmd, prefix="CONT"
+        )
+        start = time.time()
+        try:
+            subprocess.run(podman_cmd, shell=True, check=True)
+        except subprocess.CalledProcessError as e:
+            raise ContainerError(f"Build failed (exit {e.returncode})") from e
+        except KeyboardInterrupt:
+            log_warning_safe(
+                logger,
+                "Build interrupted by user - cleaning up...",
+                prefix="CONT",
+            )
+            # Get the container ID if possible
+            try:
+                container_id = (
+                    subprocess.check_output(
+                        "podman ps -q --filter ancestor=pcileech-fw-generator:latest",
+                        shell=True,
+                    )
+                    .decode()
+                    .strip()
+                )
+                if container_id:
+                    log_info_safe(
+                        logger,
+                        "Stopping container {container_id}",
+                        container_id=container_id,
+                        prefix="CONT",
+                    )
+                    subprocess.run(
+                        f"podman stop {container_id}", shell=True, check=False
+                    )
+            except Exception as e:
+                log_warning_safe(
+                    logger,
+                    "Failed to clean up container: {error}",
+                    error=str(e),
+                    prefix="CONT",
+                )
+
+            # Ensure VFIO cleanup
+            try:
+                from .vfio import restore_driver
+
+                if hasattr(cfg, "bdf") and cfg.bdf:
+                    log_info_safe(
+                        logger,
+                        "Ensuring VFIO cleanup for device {bdf}",
+                        bdf=cfg.bdf,
+                        prefix="CLEA",
+                    )
+                    # Get original driver if possible
+                    try:
+                        from .vfio import get_current_driver
+
+                        original_driver = get_current_driver(cfg.bdf)
+                        restore_driver(cfg.bdf, original_driver)
+                    except Exception:
+                        # Just try to unbind from vfio-pci
+                        try:
+                            with open(
+                                f"/sys/bus/pci/drivers/vfio-pci/unbind", "w"
+                            ) as f:
+                                f.write(f"{cfg.bdf}\n")
+                        except Exception:
+                            pass
+            except Exception as e:
+                log_warning_safe(
+                    logger,
+                    "VFIO cleanup after interrupt failed: {error}",
+                    error=str(e),
+                    prefix="CLEA",
+                )
+
+            raise KeyboardInterrupt("Build interrupted by user")
+        duration = time.time() - start
+        log_info_safe(
+            logger,
+            "Build completed in {duration:.1f}s",
+            duration=duration,
+            prefix="CONT",
+        )
+    except RuntimeError as e:
+        if "VFIO" in str(e):
+            # VFIO binding failed, diagnostics have already been run
+            log_error_safe(
+                logger,
+                "Build aborted due to VFIO issues: {error}",
+                error=str(e),
+                prefix="VFIO",
+            )
+            from .vfio_diagnostics import Diagnostics, render
+
+            # Run diagnostics one more time to ensure user sees the report
+            diag = Diagnostics(cfg.bdf)
+            report = diag.run()
+            if not report.can_proceed:
+                log_error_safe(
+                    logger,
+                    "VFIO diagnostics indicate system is not ready for VFIO operations",
+                    prefix="VFIO",
+                )
+                log_error_safe(
+                    logger,
+                    "Please fix the issues reported above and try again",
+                    prefix="VFIO",
+                )
+            sys.exit(1)
+        else:
+            # Re-raise other runtime errors
+            raise
+
+
+# ──────────────────────────────────────────────────────────────────────────────
+# CLI entry for humans / CI
+# ──────────────────────────────────────────────────────────────────────────────
+if __name__ == "__main__":
+    import argparse
+
+    p = argparse.ArgumentParser(description="VFIO‑aware Podman build wrapper")
+    p.add_argument("bdf", help="PCIe BDF, e.g. 0000:03:00.0")
+    p.add_argument("board", help="Target board string")
+    p.add_argument("--advanced-sv", action="store_true")
+    p.add_argument("--device-type", default="generic")
+    p.add_argument("--enable-variance", action="store_true")
+    p.add_argument("--auto-fix", action="store_true")
+    args = p.parse_args()
+
+    cfg = BuildConfig(
+        bdf=args.bdf,
+        board=args.board,
+        advanced_sv=args.advanced_sv,
+        device_type=args.device_type,
+        enable_variance=args.enable_variance,
+        auto_fix=args.auto_fix,
+    )
+
+    run_build(cfg)