1511 lines
57 KiB
Python
1511 lines
57 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Host-side orchestrator for DMA firmware generation.
|
|
|
|
This script:
|
|
• Enumerates PCIe devices
|
|
• Allows user to select a donor device
|
|
• Re-binds donor to vfio-pci driver
|
|
• Launches Podman container (image: pcileech-fw-generator) that runs build.py
|
|
• Optionally flashes output/firmware.bin with usbloader
|
|
• Restores original driver afterwards
|
|
|
|
Requires root privileges (sudo) for driver rebinding and VFIO operations.
|
|
"""
|
|
|
|
import argparse
|
|
import datetime
|
|
import logging
|
|
import os
|
|
import pathlib
|
|
import platform
|
|
import re
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import textwrap
|
|
import time
|
|
from typing import Dict, List, Optional, Tuple
|
|
|
|
# Import donor dump manager
|
|
try:
|
|
from src.donor_dump_manager import DonorDumpError, DonorDumpManager
|
|
except ImportError:
|
|
DonorDumpManager = None
|
|
DonorDumpError = Exception
|
|
|
|
# Git repository information
|
|
PCILEECH_FPGA_REPO = "https://github.com/ufrisk/pcileech-fpga.git"
|
|
REPO_CACHE_DIR = os.path.expanduser("~/.cache/pcileech-fw-generator/repos")
|
|
|
|
|
|
def clear_python_cache():
|
|
"""Clear Python bytecode cache files to ensure updated code is used."""
|
|
import glob
|
|
|
|
cache_patterns = [
|
|
"__pycache__",
|
|
"src/__pycache__",
|
|
"tests/__pycache__",
|
|
"tests/tui/__pycache__",
|
|
"src/tui/__pycache__",
|
|
"src/tui/core/__pycache__",
|
|
"src/tui/models/__pycache__",
|
|
"src/scripts/__pycache__",
|
|
]
|
|
|
|
for pattern in cache_patterns:
|
|
for cache_dir in glob.glob(pattern):
|
|
if os.path.exists(cache_dir):
|
|
try:
|
|
shutil.rmtree(cache_dir)
|
|
print(f"[*] Cleared Python cache: {cache_dir}")
|
|
except Exception as e:
|
|
print(f"[!] Warning: Could not clear cache {cache_dir}: {e}")
|
|
|
|
|
|
# Clear Python cache at startup to ensure updated code is used
|
|
clear_python_cache()
|
|
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format="%(asctime)s - %(levelname)s - %(message)s",
|
|
handlers=[
|
|
logging.StreamHandler(sys.stdout),
|
|
logging.FileHandler("generate.log", mode="a"),
|
|
],
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def validate_bdf_format(bdf: str) -> bool:
|
|
"""
|
|
Validate BDF (Bus:Device.Function) format.
|
|
Expected format: DDDD:BB:DD.F where D=hex digit, B=hex digit, F=0-7
|
|
Example: 0000:03:00.0
|
|
"""
|
|
bdf_pattern = re.compile(r"^[0-9a-fA-F]{4}:[0-9a-fA-F]{2}:[0-9a-fA-F]{2}\.[0-7]$")
|
|
return bool(bdf_pattern.match(bdf))
|
|
|
|
|
|
def run_command(cmd: str, timeout: int = 30, **kwargs) -> str:
|
|
"""Execute a shell command and return stripped output with timeout and better error handling."""
|
|
try:
|
|
return subprocess.check_output(
|
|
cmd, shell=True, text=True, timeout=timeout, **kwargs
|
|
).strip()
|
|
except subprocess.TimeoutExpired as e:
|
|
logger.error(f"Command timed out after {timeout}s: {cmd}")
|
|
raise RuntimeError(f"Command timed out: {cmd}") from e
|
|
except subprocess.CalledProcessError as e:
|
|
logger.error(f"Command failed (exit code {e.returncode}): {cmd}")
|
|
logger.error(f"Command stderr: {e.stderr}")
|
|
raise
|
|
|
|
|
|
def is_linux() -> bool:
|
|
"""Check if running on Linux."""
|
|
import platform
|
|
|
|
return platform.system().lower() == "linux"
|
|
|
|
|
|
def check_linux_requirement(operation: str) -> None:
|
|
"""Check if operation requires Linux and raise error if not available."""
|
|
if not is_linux():
|
|
raise RuntimeError(
|
|
f"{operation} requires Linux. "
|
|
f"Current platform: {platform.system()}. "
|
|
f"Please run this on a Linux system with VFIO support."
|
|
)
|
|
|
|
|
|
def list_pci_devices() -> List[Dict[str, str]]:
|
|
"""List all PCIe devices with their details."""
|
|
check_linux_requirement("PCIe device enumeration")
|
|
|
|
pattern = 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})\]"
|
|
)
|
|
|
|
devices = []
|
|
for line in run_command("lspci -Dnn").splitlines():
|
|
match = pattern.match(line)
|
|
if match:
|
|
device_info = match.groupdict()
|
|
device_info["pretty"] = line
|
|
devices.append(device_info)
|
|
|
|
return devices
|
|
|
|
|
|
def choose_device(devices: List[Dict[str, str]]) -> Dict[str, str]:
|
|
"""Interactive device selection from the list of PCIe devices."""
|
|
print("\nSelect donor PCIe device:")
|
|
for i, device in enumerate(devices):
|
|
print(f" [{i}] {device['pretty']}")
|
|
|
|
while True:
|
|
try:
|
|
selection = input("Enter number: ")
|
|
index = int(selection)
|
|
return devices[index]
|
|
except (ValueError, IndexError):
|
|
print(" Invalid selection — please try again.")
|
|
|
|
|
|
def get_current_driver(bdf: str) -> Optional[str]:
|
|
"""Get the current driver bound to a PCIe device."""
|
|
check_linux_requirement("Driver detection")
|
|
|
|
if not validate_bdf_format(bdf):
|
|
raise ValueError(
|
|
f"Invalid BDF format: {bdf}. Expected format: DDDD:BB:DD.F (e.g., 0000:03:00.0)"
|
|
)
|
|
|
|
driver_path = f"/sys/bus/pci/devices/{bdf}/driver"
|
|
if os.path.exists(driver_path):
|
|
return os.path.basename(os.path.realpath(driver_path))
|
|
return None
|
|
|
|
|
|
def get_iommu_group(bdf: str) -> str:
|
|
"""Get the IOMMU group for a PCIe device."""
|
|
check_linux_requirement("IOMMU group detection")
|
|
|
|
if not validate_bdf_format(bdf):
|
|
raise ValueError(
|
|
f"Invalid BDF format: {bdf}. Expected format: DDDD:BB:DD.F (e.g., 0000:03:00.0)"
|
|
)
|
|
|
|
iommu_link = f"/sys/bus/pci/devices/{bdf}/iommu_group"
|
|
return os.path.basename(os.path.realpath(iommu_link))
|
|
|
|
|
|
def list_usb_devices() -> List[Tuple[str, str]]:
|
|
"""Return list of USB devices as (vid:pid, description) tuples."""
|
|
try:
|
|
output = subprocess.check_output("lsusb", shell=True, text=True).splitlines()
|
|
except subprocess.CalledProcessError:
|
|
return []
|
|
|
|
devices = []
|
|
for line in output:
|
|
match = re.match(
|
|
r"Bus.*Device.*ID\s+([0-9a-fA-F]{4}:[0-9a-fA-F]{4})\s+(.*)", line
|
|
)
|
|
if match:
|
|
vid_pid = match.group(1)
|
|
description = match.group(2).strip() # Strip whitespace from description
|
|
devices.append((vid_pid, description))
|
|
|
|
return devices
|
|
|
|
|
|
def select_usb_device() -> str:
|
|
"""Interactive USB device selection for flashing."""
|
|
devices = list_usb_devices()
|
|
if not devices:
|
|
error_msg = "No USB devices found"
|
|
logger.error(error_msg)
|
|
raise RuntimeError(error_msg)
|
|
|
|
print("\nSelect FPGA board / USB programmer:")
|
|
for i, (vid_pid, description) in enumerate(devices):
|
|
print(f" [{i}] {vid_pid} {description}")
|
|
|
|
while True:
|
|
try:
|
|
selection = input("Enter number: ")
|
|
index = int(selection)
|
|
return devices[index][0] # Return VID:PID
|
|
except (ValueError, IndexError):
|
|
print(" Invalid selection — please try again.")
|
|
except KeyboardInterrupt:
|
|
logger.warning("USB device selection interrupted by user")
|
|
raise
|
|
|
|
|
|
def flash_firmware(bitfile: pathlib.Path) -> None:
|
|
"""Flash firmware to FPGA board using usbloader."""
|
|
logger.info("Starting firmware flash process")
|
|
|
|
if shutil.which("usbloader") is None:
|
|
error_msg = "usbloader not found in PATH — install λConcept usbloader first"
|
|
logger.error(error_msg)
|
|
raise RuntimeError(error_msg)
|
|
|
|
try:
|
|
vid_pid = select_usb_device()
|
|
logger.info(f"Selected USB device: {vid_pid}")
|
|
print(f"[*] Flashing firmware using VID:PID {vid_pid}")
|
|
|
|
subprocess.run(
|
|
f"usbloader --vidpid {vid_pid} -f {bitfile}", shell=True, check=True
|
|
)
|
|
logger.info("Firmware flashed successfully")
|
|
print("[✓] Firmware flashed successfully")
|
|
|
|
except subprocess.CalledProcessError as e:
|
|
error_msg = f"Flash failed: {e}"
|
|
logger.error(error_msg)
|
|
raise RuntimeError(error_msg) from e
|
|
except Exception as e:
|
|
error_msg = f"Unexpected error during firmware flash: {e}"
|
|
logger.error(error_msg)
|
|
raise RuntimeError(error_msg) from e
|
|
|
|
|
|
def _validate_vfio_prerequisites() -> None:
|
|
"""Validate VFIO prerequisites and system configuration."""
|
|
# Check if VFIO modules are loaded
|
|
vfio_modules = ["/sys/module/vfio", "/sys/module/vfio_pci"]
|
|
loaded_modules = [mod for mod in vfio_modules if os.path.exists(mod)]
|
|
|
|
if not loaded_modules:
|
|
error_msg = (
|
|
"VFIO modules not loaded. Please load VFIO modules:\n"
|
|
" sudo modprobe vfio\n"
|
|
" sudo modprobe vfio-pci"
|
|
)
|
|
logger.error(error_msg)
|
|
raise RuntimeError(error_msg)
|
|
|
|
# Check if vfio-pci driver is available
|
|
if not os.path.exists("/sys/bus/pci/drivers/vfio-pci"):
|
|
error_msg = (
|
|
"vfio-pci driver not available. Ensure VFIO is enabled in kernel.\n"
|
|
"Check: cat /boot/config-$(uname -r) | grep -i vfio"
|
|
)
|
|
logger.error(error_msg)
|
|
raise RuntimeError(error_msg)
|
|
|
|
# Check if IOMMU is enabled
|
|
try:
|
|
dmesg_output = run_command("dmesg | grep -i iommu", timeout=10)
|
|
if (
|
|
"IOMMU enabled" not in dmesg_output
|
|
and "AMD-Vi" not in dmesg_output
|
|
and "Intel-IOMMU" not in dmesg_output
|
|
):
|
|
logger.warning(
|
|
"IOMMU may not be enabled. Check BIOS settings and kernel parameters."
|
|
)
|
|
except Exception:
|
|
logger.debug("Could not check IOMMU status from dmesg")
|
|
|
|
|
|
def _check_device_in_use(bdf: str) -> bool:
|
|
"""Check if device is currently in use by checking for open file descriptors."""
|
|
try:
|
|
# Check if device has any open file descriptors
|
|
lsof_output = run_command(
|
|
f"lsof /dev/vfio/* 2>/dev/null | grep -v COMMAND || true", timeout=5
|
|
)
|
|
if bdf in lsof_output:
|
|
logger.warning(f"Device {bdf} may be in use by another process")
|
|
return True
|
|
except Exception:
|
|
logger.debug("Could not check device usage with lsof")
|
|
return False
|
|
|
|
|
|
def _wait_for_device_state(
|
|
bdf: str, expected_driver: Optional[str], max_retries: int = 5
|
|
) -> bool:
|
|
"""Wait for device to reach expected driver state with retries."""
|
|
import time
|
|
|
|
for attempt in range(max_retries):
|
|
try:
|
|
current_driver = get_current_driver(bdf)
|
|
if current_driver == expected_driver:
|
|
return True
|
|
|
|
if attempt < max_retries - 1:
|
|
logger.debug(
|
|
f"Device {bdf} not in expected state (current: {current_driver}, expected: {expected_driver}), retrying in 1s..."
|
|
)
|
|
time.sleep(1)
|
|
except Exception as e:
|
|
logger.debug(f"Error checking device state (attempt {attempt + 1}): {e}")
|
|
if attempt < max_retries - 1:
|
|
time.sleep(1)
|
|
|
|
return False
|
|
|
|
|
|
def bind_to_vfio(
|
|
bdf: str, vendor: str, device: str, original_driver: Optional[str]
|
|
) -> None:
|
|
"""Bind PCIe device to vfio-pci driver with enhanced error handling and validation."""
|
|
check_linux_requirement("VFIO device binding")
|
|
|
|
if not validate_bdf_format(bdf):
|
|
raise ValueError(
|
|
f"Invalid BDF format: {bdf}. Expected format: DDDD:BB:DD.F (e.g., 0000:03:00.0)"
|
|
)
|
|
|
|
# Validate vendor and device IDs
|
|
if not re.match(r"^[0-9a-fA-F]{4}$", vendor):
|
|
raise ValueError(f"Invalid vendor ID format: {vendor}. Expected 4-digit hex.")
|
|
if not re.match(r"^[0-9a-fA-F]{4}$", device):
|
|
raise ValueError(f"Invalid device ID format: {device}. Expected 4-digit hex.")
|
|
|
|
logger.info(
|
|
f"Binding device {bdf} (vendor:{vendor} device:{device}) to vfio-pci driver (current driver: {original_driver or 'none'})"
|
|
)
|
|
|
|
# Early exit if already bound to vfio-pci
|
|
if original_driver == "vfio-pci":
|
|
print(
|
|
"[*] Device already bound to vfio-pci driver, skipping binding process..."
|
|
)
|
|
logger.info(f"Device {bdf} already bound to vfio-pci, skipping binding process")
|
|
return
|
|
|
|
print("[*] Binding device to vfio-pci driver...")
|
|
|
|
try:
|
|
# Validate VFIO prerequisites
|
|
_validate_vfio_prerequisites()
|
|
|
|
# Check if device exists
|
|
device_path = f"/sys/bus/pci/devices/{bdf}"
|
|
if not os.path.exists(device_path):
|
|
error_msg = f"PCIe device {bdf} not found in sysfs"
|
|
logger.error(error_msg)
|
|
raise RuntimeError(error_msg)
|
|
|
|
# Check if device is in use
|
|
if _check_device_in_use(bdf):
|
|
logger.warning(
|
|
f"Device {bdf} appears to be in use, proceeding with caution"
|
|
)
|
|
|
|
# Check if device ID is already registered with vfio-pci
|
|
device_id_registered = False
|
|
try:
|
|
ids_file = "/sys/bus/pci/drivers/vfio-pci/ids"
|
|
if os.path.exists(ids_file):
|
|
with open(ids_file, "r") as f:
|
|
registered_ids = f.read()
|
|
if f"{vendor} {device}" in registered_ids:
|
|
logger.info(
|
|
f"Device ID {vendor}:{device} already registered with vfio-pci"
|
|
)
|
|
device_id_registered = True
|
|
except (OSError, IOError) as e:
|
|
logger.debug(f"Error checking registered device IDs: {e}")
|
|
# Continue with normal flow if we can't check
|
|
|
|
# Register device ID with vfio-pci if not already registered
|
|
if not device_id_registered:
|
|
logger.debug(f"Registering device ID {vendor}:{device} with vfio-pci")
|
|
max_retries = 3
|
|
for attempt in range(max_retries):
|
|
try:
|
|
# Use direct file writing instead of shell redirection to avoid I/O errors
|
|
new_id_path = "/sys/bus/pci/drivers/vfio-pci/new_id"
|
|
with open(new_id_path, "w") as f:
|
|
f.write(f"{vendor} {device}\n")
|
|
logger.info(
|
|
f"Successfully registered device ID {vendor}:{device} with vfio-pci"
|
|
)
|
|
break
|
|
except (OSError, IOError) as e:
|
|
if (
|
|
"File exists" in str(e)
|
|
or "Invalid argument" in str(e)
|
|
or "Device or resource busy" in str(e)
|
|
):
|
|
logger.info(
|
|
f"Device ID {vendor}:{device} already registered with vfio-pci or busy"
|
|
)
|
|
break
|
|
elif attempt < max_retries - 1:
|
|
logger.warning(
|
|
f"Failed to register device ID (attempt {attempt + 1}): {e}, retrying..."
|
|
)
|
|
import time
|
|
|
|
time.sleep(1)
|
|
else:
|
|
logger.error(
|
|
f"Failed to register device ID after {max_retries} attempts: {e}"
|
|
)
|
|
raise RuntimeError(f"Failed to register device ID: {e}")
|
|
|
|
# Unbind from current driver if present
|
|
if original_driver:
|
|
logger.debug(f"Unbinding from current driver: {original_driver}")
|
|
max_retries = 3
|
|
for attempt in range(max_retries):
|
|
try:
|
|
# Use direct file writing instead of shell redirection
|
|
unbind_path = f"/sys/bus/pci/devices/{bdf}/driver/unbind"
|
|
with open(unbind_path, "w") as f:
|
|
f.write(f"{bdf}\n")
|
|
logger.info(f"Successfully unbound {bdf} from {original_driver}")
|
|
|
|
# Wait for unbind to complete
|
|
if _wait_for_device_state(bdf, None, max_retries=3):
|
|
break
|
|
elif attempt < max_retries - 1:
|
|
logger.warning(
|
|
f"Device still bound after unbind (attempt {attempt + 1}), retrying..."
|
|
)
|
|
import time
|
|
|
|
time.sleep(1)
|
|
else:
|
|
logger.warning(
|
|
"Device may still be bound to original driver, continuing..."
|
|
)
|
|
break
|
|
|
|
except (OSError, IOError) as e:
|
|
if "No such device" in str(e) or "No such file or directory" in str(
|
|
e
|
|
):
|
|
logger.info(f"Device {bdf} already unbound")
|
|
break
|
|
elif attempt < max_retries - 1:
|
|
logger.warning(
|
|
f"Failed to unbind from current driver (attempt {attempt + 1}): {e}, retrying..."
|
|
)
|
|
import time
|
|
|
|
time.sleep(1)
|
|
else:
|
|
logger.warning(
|
|
f"Failed to unbind from current driver after {max_retries} attempts: {e}"
|
|
)
|
|
# Continue anyway, as the bind might still work
|
|
|
|
# Bind to vfio-pci with retries
|
|
logger.debug(f"Binding {bdf} to vfio-pci")
|
|
max_retries = 3
|
|
bind_successful = False
|
|
|
|
for attempt in range(max_retries):
|
|
try:
|
|
# Use direct file writing instead of shell redirection
|
|
bind_path = "/sys/bus/pci/drivers/vfio-pci/bind"
|
|
with open(bind_path, "w") as f:
|
|
f.write(f"{bdf}\n")
|
|
|
|
# Verify binding was successful
|
|
if _wait_for_device_state(bdf, "vfio-pci", max_retries=3):
|
|
logger.info(f"Successfully bound {bdf} to vfio-pci")
|
|
print("[✓] Device successfully bound to vfio-pci driver")
|
|
bind_successful = True
|
|
break
|
|
elif attempt < max_retries - 1:
|
|
logger.warning(
|
|
f"Bind command succeeded but device not bound to vfio-pci (attempt {attempt + 1}), retrying..."
|
|
)
|
|
import time
|
|
|
|
time.sleep(1)
|
|
|
|
except (OSError, IOError) as e:
|
|
if "Device or resource busy" in str(e):
|
|
logger.warning(f"Device {bdf} is busy (attempt {attempt + 1})")
|
|
if attempt < max_retries - 1:
|
|
import time
|
|
|
|
time.sleep(2) # Longer wait for busy devices
|
|
continue
|
|
elif "No such device" in str(e) or "No such file or directory" in str(
|
|
e
|
|
):
|
|
logger.error(f"Device {bdf} disappeared during binding")
|
|
raise RuntimeError(f"Device {bdf} not found during binding")
|
|
elif attempt < max_retries - 1:
|
|
logger.warning(
|
|
f"Failed to bind to vfio-pci (attempt {attempt + 1}): {e}, retrying..."
|
|
)
|
|
import time
|
|
|
|
time.sleep(1)
|
|
else:
|
|
# Final attempt failed, check if device is actually bound
|
|
current_driver = get_current_driver(bdf)
|
|
if current_driver == "vfio-pci":
|
|
logger.info(
|
|
f"Device {bdf} is bound to vfio-pci despite bind command error"
|
|
)
|
|
print("[✓] Device is bound to vfio-pci driver")
|
|
bind_successful = True
|
|
break
|
|
else:
|
|
logger.error(
|
|
f"Failed to bind to vfio-pci after {max_retries} attempts: {e}"
|
|
)
|
|
raise RuntimeError(f"Failed to bind to vfio-pci: {e}")
|
|
|
|
if not bind_successful:
|
|
error_msg = (
|
|
f"Failed to bind device {bdf} to vfio-pci after {max_retries} attempts"
|
|
)
|
|
logger.error(error_msg)
|
|
raise RuntimeError(error_msg)
|
|
|
|
# Final verification
|
|
final_driver = get_current_driver(bdf)
|
|
if final_driver != "vfio-pci":
|
|
error_msg = (
|
|
f"Device {bdf} not bound to vfio-pci (current driver: {final_driver})"
|
|
)
|
|
logger.error(error_msg)
|
|
raise RuntimeError(error_msg)
|
|
|
|
except (OSError, IOError) as e:
|
|
error_msg = f"Failed to bind device to vfio-pci: {e}"
|
|
logger.error(error_msg)
|
|
raise RuntimeError(error_msg) from e
|
|
except Exception as e:
|
|
error_msg = f"Unexpected error during vfio binding: {e}"
|
|
logger.error(error_msg)
|
|
raise RuntimeError(error_msg) from e
|
|
|
|
|
|
def restore_original_driver(bdf: str, original_driver: Optional[str]) -> None:
|
|
"""Restore the original driver binding for the PCIe device with enhanced error handling."""
|
|
if not validate_bdf_format(bdf):
|
|
logger.warning(f"Invalid BDF format during restore: {bdf}")
|
|
return
|
|
|
|
logger.info(
|
|
f"Restoring original driver binding for {bdf} (target driver: {original_driver or 'none'})"
|
|
)
|
|
print("[*] Restoring original driver binding...")
|
|
|
|
try:
|
|
# Check if device exists
|
|
device_path = f"/sys/bus/pci/devices/{bdf}"
|
|
if not os.path.exists(device_path):
|
|
logger.warning(
|
|
f"Device {bdf} not found during restore, may have been removed"
|
|
)
|
|
return
|
|
|
|
# Check current driver state
|
|
current_driver = get_current_driver(bdf)
|
|
logger.debug(f"Current driver for {bdf}: {current_driver or 'none'}")
|
|
|
|
# Unbind from vfio-pci if currently bound
|
|
if current_driver == "vfio-pci":
|
|
logger.debug(f"Unbinding {bdf} from vfio-pci")
|
|
max_retries = 3
|
|
unbind_successful = False
|
|
|
|
for attempt in range(max_retries):
|
|
try:
|
|
# Use direct file writing instead of shell redirection
|
|
unbind_path = "/sys/bus/pci/drivers/vfio-pci/unbind"
|
|
with open(unbind_path, "w") as f:
|
|
f.write(f"{bdf}\n")
|
|
|
|
# Wait for unbind to complete
|
|
if _wait_for_device_state(bdf, None, max_retries=3):
|
|
logger.info(f"Successfully unbound {bdf} from vfio-pci")
|
|
unbind_successful = True
|
|
break
|
|
elif attempt < max_retries - 1:
|
|
logger.warning(
|
|
f"Device still bound to vfio-pci (attempt {attempt + 1}), retrying..."
|
|
)
|
|
import time
|
|
|
|
time.sleep(1)
|
|
|
|
except (OSError, IOError) as e:
|
|
if "No such device" in str(e) or "No such file or directory" in str(
|
|
e
|
|
):
|
|
logger.info(f"Device {bdf} already unbound from vfio-pci")
|
|
unbind_successful = True
|
|
break
|
|
elif attempt < max_retries - 1:
|
|
logger.warning(
|
|
f"Failed to unbind from vfio-pci (attempt {attempt + 1}): {e}, retrying..."
|
|
)
|
|
import time
|
|
|
|
time.sleep(1)
|
|
else:
|
|
logger.warning(
|
|
f"Failed to unbind from vfio-pci after {max_retries} attempts: {e}"
|
|
)
|
|
# Continue with restore attempt anyway
|
|
break
|
|
|
|
if not unbind_successful and get_current_driver(bdf) == "vfio-pci":
|
|
logger.warning(
|
|
f"Device {bdf} still bound to vfio-pci, restore may fail"
|
|
)
|
|
else:
|
|
logger.info(
|
|
f"Device {bdf} not bound to vfio-pci (current: {current_driver or 'none'})"
|
|
)
|
|
|
|
# Bind back to original driver if it existed
|
|
if original_driver:
|
|
# Check if original driver is available
|
|
driver_path = f"/sys/bus/pci/drivers/{original_driver}"
|
|
if not os.path.exists(driver_path):
|
|
logger.warning(
|
|
f"Original driver {original_driver} not available for restore"
|
|
)
|
|
print(f"Warning: Original driver {original_driver} not available")
|
|
return
|
|
|
|
logger.debug(f"Binding {bdf} back to {original_driver}")
|
|
max_retries = 3
|
|
restore_successful = False
|
|
|
|
for attempt in range(max_retries):
|
|
try:
|
|
# Use direct file writing instead of shell redirection
|
|
bind_path = f"/sys/bus/pci/drivers/{original_driver}/bind"
|
|
with open(bind_path, "w") as f:
|
|
f.write(f"{bdf}\n")
|
|
|
|
# Verify restore was successful
|
|
if _wait_for_device_state(bdf, original_driver, max_retries=3):
|
|
logger.info(f"Successfully restored {bdf} to {original_driver}")
|
|
print(f"[✓] Device restored to {original_driver} driver")
|
|
restore_successful = True
|
|
break
|
|
elif attempt < max_retries - 1:
|
|
logger.warning(
|
|
f"Restore command succeeded but device not bound to {original_driver} (attempt {attempt + 1}), retrying..."
|
|
)
|
|
import time
|
|
|
|
time.sleep(1)
|
|
|
|
except (OSError, IOError) as e:
|
|
if "Device or resource busy" in str(e):
|
|
logger.warning(
|
|
f"Device {bdf} is busy during restore (attempt {attempt + 1})"
|
|
)
|
|
if attempt < max_retries - 1:
|
|
import time
|
|
|
|
time.sleep(2)
|
|
continue
|
|
elif "No such device" in str(
|
|
e
|
|
) or "No such file or directory" in str(e):
|
|
logger.warning(f"Device {bdf} not found during restore")
|
|
break
|
|
elif attempt < max_retries - 1:
|
|
logger.warning(
|
|
f"Failed to restore to {original_driver} (attempt {attempt + 1}): {e}, retrying..."
|
|
)
|
|
import time
|
|
|
|
time.sleep(1)
|
|
else:
|
|
logger.warning(
|
|
f"Failed to restore to {original_driver} after {max_retries} attempts: {e}"
|
|
)
|
|
break
|
|
|
|
if not restore_successful:
|
|
final_driver = get_current_driver(bdf)
|
|
if final_driver == original_driver:
|
|
logger.info(
|
|
f"Device {bdf} is bound to {original_driver} despite restore errors"
|
|
)
|
|
print(f"[✓] Device is bound to {original_driver} driver")
|
|
else:
|
|
logger.warning(
|
|
f"Failed to restore {bdf} to {original_driver}, current driver: {final_driver or 'none'}"
|
|
)
|
|
print(
|
|
f"Warning: Failed to restore to {original_driver}, current driver: {final_driver or 'none'}"
|
|
)
|
|
else:
|
|
logger.info(f"No original driver to restore for {bdf}")
|
|
print("[*] No original driver to restore")
|
|
|
|
except (OSError, IOError) as e:
|
|
logger.warning(f"Failed to restore original driver for {bdf}: {e}")
|
|
print(f"Warning: Failed to restore original driver: {e}")
|
|
except Exception as e:
|
|
logger.warning(f"Unexpected error during driver restore for {bdf}: {e}")
|
|
print(f"Warning: Unexpected error during driver restore: {e}")
|
|
|
|
|
|
def _validate_vfio_device_access(vfio_device: str, bdf: str) -> None:
|
|
"""Validate VFIO device access and permissions."""
|
|
# Check if VFIO device exists
|
|
if not os.path.exists(vfio_device):
|
|
error_msg = f"VFIO device {vfio_device} not found"
|
|
logger.error(error_msg)
|
|
raise RuntimeError(error_msg)
|
|
|
|
# Check if /dev/vfio/vfio exists (VFIO container device)
|
|
vfio_container = "/dev/vfio/vfio"
|
|
if not os.path.exists(vfio_container):
|
|
error_msg = f"VFIO container device {vfio_container} not found"
|
|
logger.error(error_msg)
|
|
raise RuntimeError(error_msg)
|
|
|
|
# Check device permissions
|
|
try:
|
|
import stat
|
|
|
|
vfio_stat = os.stat(vfio_device)
|
|
if not (vfio_stat.st_mode & stat.S_IRGRP) or not (
|
|
vfio_stat.st_mode & stat.S_IWGRP
|
|
):
|
|
logger.warning(
|
|
f"VFIO device {vfio_device} may not have proper group permissions"
|
|
)
|
|
except OSError as e:
|
|
logger.warning(f"Could not check VFIO device permissions: {e}")
|
|
|
|
# Verify device is actually bound to vfio-pci
|
|
current_driver = get_current_driver(bdf)
|
|
if current_driver != "vfio-pci":
|
|
error_msg = (
|
|
f"Device {bdf} not bound to vfio-pci (current: {current_driver or 'none'})"
|
|
)
|
|
logger.error(error_msg)
|
|
raise RuntimeError(error_msg)
|
|
|
|
|
|
def _validate_container_environment() -> None:
|
|
"""Validate container runtime environment."""
|
|
# Check if podman is available
|
|
if shutil.which("podman") is None:
|
|
error_msg = "Podman not found in PATH. Please install Podman container runtime."
|
|
logger.error(error_msg)
|
|
raise RuntimeError(error_msg)
|
|
|
|
# Check if container image exists
|
|
try:
|
|
result = run_command(
|
|
"podman images --format '{{.Repository}}:{{.Tag}}' | grep '^pcileech-fw-generator:'",
|
|
timeout=10,
|
|
)
|
|
if not result:
|
|
# Container image not found, try to build it automatically
|
|
logger.info(
|
|
"Container image 'pcileech-fw-generator' not found. Building it now..."
|
|
)
|
|
print(
|
|
"[*] Container image 'pcileech-fw-generator' not found. Building it now..."
|
|
)
|
|
|
|
try:
|
|
# Use the proper build script which handles the container building correctly
|
|
build_script_path = "scripts/build_container.sh"
|
|
if os.path.exists(build_script_path):
|
|
logger.info("Using build script for container creation")
|
|
build_result = subprocess.run(
|
|
f"bash {build_script_path} --tag pcileech-fw-generator:latest",
|
|
shell=True,
|
|
check=True,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True,
|
|
)
|
|
else:
|
|
# Fallback to direct container build
|
|
build_result = subprocess.run(
|
|
"podman build -t pcileech-fw-generator:latest -f Containerfile .",
|
|
shell=True,
|
|
check=True,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True,
|
|
)
|
|
logger.info("Container image built successfully")
|
|
print("[✓] Container image built successfully")
|
|
except subprocess.CalledProcessError as e:
|
|
error_msg = f"Failed to build container image automatically: {e.stderr}\nPlease build manually with: make container"
|
|
logger.error(error_msg)
|
|
raise RuntimeError(error_msg)
|
|
except subprocess.CalledProcessError:
|
|
# If we can't check, try to build anyway
|
|
logger.info("Could not check container image status. Attempting to build...")
|
|
print("[*] Could not check container image status. Attempting to build...")
|
|
|
|
try:
|
|
build_script_path = "scripts/build_container.sh"
|
|
if os.path.exists(build_script_path):
|
|
logger.info("Using build script for container creation")
|
|
build_result = subprocess.run(
|
|
f"bash {build_script_path} --tag pcileech-fw-generator:latest",
|
|
shell=True,
|
|
check=True,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True,
|
|
)
|
|
else:
|
|
# Fallback to direct container build
|
|
build_result = subprocess.run(
|
|
"podman build -t pcileech-fw-generator:latest -f Containerfile .",
|
|
shell=True,
|
|
check=True,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True,
|
|
)
|
|
logger.info("Container image built successfully")
|
|
print("[✓] Container image built successfully")
|
|
except subprocess.CalledProcessError as e:
|
|
error_msg = f"Failed to build container image: {e.stderr}\nPlease build manually with: make container"
|
|
logger.error(error_msg)
|
|
raise RuntimeError(error_msg)
|
|
|
|
|
|
def run_build_container(
|
|
bdf: str, board: str, vfio_device: str, args: argparse.Namespace
|
|
) -> None:
|
|
"""Run the firmware build in a Podman container with enhanced validation and error handling."""
|
|
if not validate_bdf_format(bdf):
|
|
raise ValueError(
|
|
f"Invalid BDF format: {bdf}. Expected format: DDDD:BB:DD.F (e.g., 0000:03:00.0)"
|
|
)
|
|
|
|
# Log advanced features being used
|
|
advanced_features = []
|
|
if args.advanced_sv:
|
|
advanced_features.append("Advanced SystemVerilog Generation")
|
|
if args.enable_variance:
|
|
advanced_features.append("Manufacturing Variance Simulation")
|
|
if args.device_type != "generic":
|
|
advanced_features.append(f"Device-specific optimizations ({args.device_type})")
|
|
|
|
if advanced_features:
|
|
logger.info(f"Advanced features enabled: {', '.join(advanced_features)}")
|
|
print(f"[*] Advanced features: {', '.join(advanced_features)}")
|
|
|
|
logger.info(f"Starting container build for device {bdf} on board {board}")
|
|
|
|
# Validate container environment
|
|
_validate_container_environment()
|
|
|
|
# Ensure output directory exists with proper permissions
|
|
output_dir = "output"
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
|
|
# Check output directory permissions
|
|
if not os.access(output_dir, os.W_OK):
|
|
error_msg = f"Output directory {output_dir} is not writable"
|
|
logger.error(error_msg)
|
|
raise RuntimeError(error_msg)
|
|
|
|
# Validate VFIO device access
|
|
_validate_vfio_device_access(vfio_device, bdf)
|
|
|
|
# Build the build.py command with all arguments - use modular build system if available
|
|
build_cmd_parts = [f"sudo python3 /app/src/build.py --bdf {bdf} --board {board}"]
|
|
|
|
# Add advanced features arguments
|
|
if args.advanced_sv:
|
|
build_cmd_parts.append("--advanced-sv")
|
|
|
|
if args.device_type != "generic":
|
|
build_cmd_parts.append(f"--device-type {args.device_type}")
|
|
|
|
if args.enable_variance:
|
|
build_cmd_parts.append("--enable-variance")
|
|
|
|
if args.disable_power_management:
|
|
build_cmd_parts.append("--disable-power-management")
|
|
|
|
if args.disable_error_handling:
|
|
build_cmd_parts.append("--disable-error-handling")
|
|
|
|
if args.disable_performance_counters:
|
|
build_cmd_parts.append("--disable-performance-counters")
|
|
|
|
if args.behavior_profile_duration != 30:
|
|
build_cmd_parts.append(
|
|
f"--behavior-profile-duration {args.behavior_profile_duration}"
|
|
)
|
|
|
|
build_cmd = " ".join(build_cmd_parts)
|
|
|
|
# Construct Podman command
|
|
container_cmd = textwrap.dedent(
|
|
f"""
|
|
podman run --rm -it --privileged \
|
|
--device={vfio_device} \
|
|
--device=/dev/vfio/vfio \
|
|
-v {os.getcwd()}/output:/app/output \
|
|
pcileech-fw-generator:latest \
|
|
{build_cmd}
|
|
"""
|
|
).strip()
|
|
|
|
logger.debug(f"Container command: {container_cmd}")
|
|
print("[*] Launching build container...")
|
|
start_time = time.time()
|
|
|
|
try:
|
|
subprocess.run(container_cmd, shell=True, check=True)
|
|
elapsed_time = time.time() - start_time
|
|
logger.info(f"Build completed successfully in {elapsed_time:.1f} seconds")
|
|
print(f"[✓] Build completed in {elapsed_time:.1f} seconds")
|
|
|
|
except subprocess.CalledProcessError as e:
|
|
elapsed_time = time.time() - start_time
|
|
error_msg = f"Container build failed after {elapsed_time:.1f} seconds: {e}"
|
|
logger.error(error_msg)
|
|
raise RuntimeError(error_msg) from e
|
|
except Exception as e:
|
|
elapsed_time = time.time() - start_time
|
|
error_msg = f"Unexpected error during container build after {elapsed_time:.1f} seconds: {e}"
|
|
logger.error(error_msg)
|
|
raise RuntimeError(error_msg) from e
|
|
|
|
|
|
def ensure_git_repo(repo_url: str, local_dir: str, update: bool = False) -> str:
|
|
"""
|
|
Ensure that the git repository is available locally.
|
|
|
|
Args:
|
|
repo_url (str): URL of the git repository
|
|
local_dir (str): Local directory to clone/pull the repository
|
|
update (bool): Whether to update the repository if it already exists
|
|
|
|
Returns:
|
|
str: Path to the local repository
|
|
"""
|
|
# Create cache directory if it doesn't exist
|
|
os.makedirs(os.path.dirname(local_dir), exist_ok=True)
|
|
|
|
# Check if repository already exists as a valid git repository
|
|
if os.path.exists(os.path.join(local_dir, ".git")):
|
|
logger.info(f"Repository already exists at {local_dir}")
|
|
|
|
# Update repository if requested
|
|
if update:
|
|
try:
|
|
logger.info(f"Updating repository at {local_dir}")
|
|
print(f"[*] Updating repository at {local_dir}")
|
|
|
|
# Get current directory
|
|
current_dir = os.getcwd()
|
|
|
|
# Change to repository directory
|
|
os.chdir(local_dir)
|
|
|
|
# Pull latest changes
|
|
result = subprocess.run(
|
|
"git pull",
|
|
shell=True,
|
|
check=True,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True,
|
|
)
|
|
|
|
# Change back to original directory
|
|
os.chdir(current_dir)
|
|
|
|
logger.info(f"Repository updated successfully: {result.stdout.strip()}")
|
|
print(f"[✓] Repository updated successfully")
|
|
except subprocess.CalledProcessError as e:
|
|
logger.warning(f"Failed to update repository: {e.stderr}")
|
|
print(f"[!] Warning: Failed to update repository: {e.stderr}")
|
|
else:
|
|
# Check if directory exists but is not a git repository
|
|
if os.path.exists(local_dir):
|
|
logger.info(f"Directory exists but is not a git repository: {local_dir}")
|
|
print(f"[*] Removing existing directory: {local_dir}")
|
|
|
|
# Remove the directory to allow fresh clone
|
|
import shutil
|
|
|
|
try:
|
|
shutil.rmtree(local_dir)
|
|
except Exception as e:
|
|
logger.warning(f"Failed to remove directory: {e}")
|
|
print(f"[!] Warning: Failed to remove directory: {e}")
|
|
# Continue anyway, git clone might still work
|
|
|
|
# Clone repository
|
|
try:
|
|
logger.info(f"Cloning repository {repo_url} to {local_dir}")
|
|
print(f"[*] Cloning repository {repo_url} to {local_dir}")
|
|
|
|
result = subprocess.run(
|
|
f"git clone {repo_url} {local_dir}",
|
|
shell=True,
|
|
check=True,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True,
|
|
)
|
|
|
|
logger.info(f"Repository cloned successfully")
|
|
print(f"[✓] Repository cloned successfully")
|
|
except subprocess.CalledProcessError as e:
|
|
error_msg = f"Failed to clone repository: {e.stderr}"
|
|
logger.error(error_msg)
|
|
raise RuntimeError(error_msg)
|
|
|
|
# Return path to repository
|
|
return local_dir
|
|
|
|
|
|
def validate_environment() -> None:
|
|
"""Validate that the environment is properly set up."""
|
|
if os.geteuid() != 0:
|
|
error_msg = "This script requires root privileges. Run with sudo."
|
|
logger.error(error_msg)
|
|
raise RuntimeError(error_msg)
|
|
|
|
# Check if git is available
|
|
if shutil.which("git") is None:
|
|
error_msg = "Git not found in PATH. Please install Git first."
|
|
logger.error(error_msg)
|
|
raise RuntimeError(error_msg)
|
|
|
|
# Check if Podman is available
|
|
if shutil.which("podman") is None:
|
|
error_msg = "Podman not found in PATH. Please install Podman first."
|
|
logger.error(error_msg)
|
|
raise RuntimeError(error_msg)
|
|
|
|
# Check if Vivado is available
|
|
try:
|
|
# Import vivado_utils from src directory
|
|
from pathlib import Path
|
|
|
|
sys.path.insert(0, str(Path(__file__).parent / "src"))
|
|
from src.vivado_utils import find_vivado_installation, get_vivado_search_paths
|
|
|
|
vivado_info = find_vivado_installation()
|
|
if vivado_info:
|
|
logger.info(
|
|
f"Found Vivado {vivado_info['version']} at {vivado_info['path']}"
|
|
)
|
|
print(f"[✓] Vivado {vivado_info['version']} detected")
|
|
else:
|
|
# Show what paths were checked using the utility function
|
|
checked_paths = get_vivado_search_paths()
|
|
error_msg = f"Vivado not found. Checked paths:\n" + "\n".join(
|
|
f" • {path}" for path in checked_paths
|
|
)
|
|
logger.error(error_msg)
|
|
print(f"[✗] {error_msg}")
|
|
|
|
# Allow user to skip
|
|
try:
|
|
response = (
|
|
input("\nWould you like to continue without Vivado? (y/N): ")
|
|
.strip()
|
|
.lower()
|
|
)
|
|
if response in ["y", "yes"]:
|
|
print("[!] Continuing without Vivado - some features may not work")
|
|
logger.warning("User chose to continue without Vivado")
|
|
else:
|
|
raise RuntimeError(
|
|
"Vivado is required. Please install Vivado and try again."
|
|
)
|
|
except (KeyboardInterrupt, EOFError):
|
|
raise RuntimeError(
|
|
"Vivado is required. Please install Vivado and try again."
|
|
)
|
|
except ImportError:
|
|
error_msg = (
|
|
"Could not import vivado_utils. Please ensure Vivado is properly installed."
|
|
)
|
|
logger.error(error_msg)
|
|
print(f"[✗] {error_msg}")
|
|
|
|
# Allow user to skip
|
|
try:
|
|
response = (
|
|
input("\nWould you like to continue without Vivado? (y/N): ")
|
|
.strip()
|
|
.lower()
|
|
)
|
|
if response in ["y", "yes"]:
|
|
print("[!] Continuing without Vivado - some features may not work")
|
|
logger.warning("User chose to continue without Vivado")
|
|
else:
|
|
raise RuntimeError(
|
|
"Vivado is required. Please install Vivado and try again."
|
|
)
|
|
except (KeyboardInterrupt, EOFError):
|
|
raise RuntimeError(
|
|
"Vivado is required. Please install Vivado and try again."
|
|
)
|
|
|
|
# Check if container image exists
|
|
try:
|
|
result = run_command(
|
|
"podman images pcileech-fw-generator --format '{{.Repository}}'"
|
|
)
|
|
if "pcileech-fw-generator" not in result:
|
|
# Container image not found, try to build it
|
|
logger.info(
|
|
"Container image 'pcileech-fw-generator' not found. Building it now..."
|
|
)
|
|
print(
|
|
"[*] Container image 'pcileech-fw-generator' not found. Building it now..."
|
|
)
|
|
|
|
try:
|
|
# Use the proper build script which handles the container building correctly
|
|
build_script_path = "scripts/build_container.sh"
|
|
if os.path.exists(build_script_path):
|
|
logger.info("Using build script for container creation")
|
|
build_result = subprocess.run(
|
|
f"bash {build_script_path} --tag pcileech-fw-generator:latest",
|
|
shell=True,
|
|
check=True,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True,
|
|
)
|
|
else:
|
|
# Fallback to direct container build
|
|
build_result = subprocess.run(
|
|
"podman build -t pcileech-fw-generator:latest -f Containerfile .",
|
|
shell=True,
|
|
check=True,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True,
|
|
)
|
|
logger.info("Container image built successfully")
|
|
print("[✓] Container image built successfully")
|
|
except subprocess.CalledProcessError as e:
|
|
error_msg = f"Failed to build container image: {e.stderr}"
|
|
logger.error(error_msg)
|
|
raise RuntimeError(error_msg)
|
|
except Exception as e:
|
|
error_msg = f"Error checking container image: {str(e)}"
|
|
logger.error(error_msg)
|
|
raise RuntimeError(error_msg)
|
|
|
|
|
|
def main() -> int:
|
|
"""Main entry point for the firmware generator"""
|
|
bdf = None
|
|
original_driver = None
|
|
|
|
try:
|
|
logger.info("Starting PCILeech firmware generation process")
|
|
validate_environment()
|
|
|
|
# Ensure pcileech-fpga repository is available
|
|
repo_dir = os.path.join(REPO_CACHE_DIR, "pcileech-fpga")
|
|
pcileech_fpga_dir = ensure_git_repo(PCILEECH_FPGA_REPO, repo_dir, update=False)
|
|
logger.info(f"Using pcileech-fpga repository at {pcileech_fpga_dir}")
|
|
|
|
# Parse command line arguments
|
|
parser = argparse.ArgumentParser(
|
|
description="Generate DMA firmware from donor PCIe device",
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
epilog=textwrap.dedent(
|
|
"""
|
|
Examples:
|
|
# Basic usage
|
|
sudo python3 generate.py --board 75t
|
|
sudo python3 generate.py --board 100t --flash
|
|
|
|
# Advanced SystemVerilog generation
|
|
sudo python3 generate.py --board 75t --advanced-sv --device-type network
|
|
|
|
# Manufacturing variance simulation
|
|
sudo python3 generate.py --board 100t --enable-variance --behavior-profile-duration 60
|
|
|
|
# Advanced features with selective disabling
|
|
sudo python3 generate.py --board 75t --advanced-sv --disable-power-management --disable-error-handling
|
|
|
|
# Full advanced configuration
|
|
sudo python3 generate.py --board 100t --advanced-sv --device-type storage --enable-variance --behavior-profile-duration 45 --flash
|
|
"""
|
|
),
|
|
)
|
|
|
|
# Basic options
|
|
parser.add_argument(
|
|
"--tui",
|
|
action="store_true",
|
|
help="Launch TUI (Text User Interface) mode",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--flash",
|
|
action="store_true",
|
|
help="Flash output/firmware.bin with usbloader after build",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--board",
|
|
choices=[
|
|
# Original boards
|
|
"35t",
|
|
"75t",
|
|
"100t",
|
|
# CaptainDMA boards
|
|
"pcileech_75t484_x1",
|
|
"pcileech_35t484_x1",
|
|
"pcileech_35t325_x4",
|
|
"pcileech_35t325_x1",
|
|
"pcileech_100t484_x1",
|
|
# Other boards
|
|
"pcileech_enigma_x1",
|
|
"pcileech_squirrel",
|
|
"pcileech_pciescreamer_xc7a35",
|
|
],
|
|
default="35t",
|
|
help="Target FPGA board type (default: 35t/Squirrel)",
|
|
)
|
|
|
|
# Advanced SystemVerilog Generation
|
|
parser.add_argument(
|
|
"--advanced-sv",
|
|
action="store_true",
|
|
help="Enable advanced SystemVerilog generation with enhanced features",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--device-type",
|
|
choices=["network", "storage", "graphics", "audio", "generic"],
|
|
default="generic",
|
|
help="Device type for specialized optimizations (default: generic)",
|
|
)
|
|
|
|
# Manufacturing Variance Simulation
|
|
parser.add_argument(
|
|
"--enable-variance",
|
|
action="store_true",
|
|
help="Enable manufacturing variance simulation for realistic timing",
|
|
)
|
|
|
|
# Feature Control
|
|
parser.add_argument(
|
|
"--disable-power-management",
|
|
action="store_true",
|
|
help="Disable power management features in advanced generation",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--disable-error-handling",
|
|
action="store_true",
|
|
help="Disable error handling features in advanced generation",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--disable-performance-counters",
|
|
action="store_true",
|
|
help="Disable performance counter features in advanced generation",
|
|
)
|
|
|
|
# Behavior Profiling
|
|
parser.add_argument(
|
|
"--behavior-profile-duration",
|
|
type=int,
|
|
default=30,
|
|
help="Duration for behavior profiling in seconds (default: 30)",
|
|
)
|
|
|
|
# Donor dump functionality
|
|
parser.add_argument(
|
|
"--donor-dump",
|
|
action="store_true",
|
|
help="Extract donor device parameters using kernel module before generation",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--auto-install-headers",
|
|
action="store_true",
|
|
help="Automatically install kernel headers if missing (for donor dump)",
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Check if TUI mode is requested
|
|
if args.tui:
|
|
try:
|
|
# Import and launch TUI
|
|
from src.tui.main import PCILeechTUI
|
|
|
|
app = PCILeechTUI()
|
|
app.run()
|
|
return 0
|
|
except ImportError:
|
|
error_msg = "TUI dependencies not installed. Install with: pip install textual rich psutil"
|
|
logger.error(error_msg)
|
|
print(f"[✗] {error_msg}")
|
|
return 1
|
|
except Exception as e:
|
|
error_msg = f"Failed to launch TUI: {e}"
|
|
logger.error(error_msg)
|
|
print(f"[✗] {error_msg}")
|
|
return 1
|
|
|
|
# Enhanced logging with advanced features
|
|
config_info = [f"board={args.board}", f"flash={args.flash}"]
|
|
if args.advanced_sv:
|
|
config_info.append("advanced_sv=True")
|
|
if args.device_type != "generic":
|
|
config_info.append(f"device_type={args.device_type}")
|
|
if args.enable_variance:
|
|
config_info.append("variance=True")
|
|
if args.disable_power_management:
|
|
config_info.append("no_power_mgmt=True")
|
|
if args.disable_error_handling:
|
|
config_info.append("no_error_handling=True")
|
|
if args.disable_performance_counters:
|
|
config_info.append("no_perf_counters=True")
|
|
if args.behavior_profile_duration != 30:
|
|
config_info.append(f"profile_duration={args.behavior_profile_duration}s")
|
|
|
|
logger.info(f"Configuration: {', '.join(config_info)}")
|
|
|
|
# List and select PCIe device
|
|
devices = list_pci_devices()
|
|
if not devices:
|
|
error_msg = "No PCIe devices found"
|
|
logger.error(error_msg)
|
|
raise RuntimeError(error_msg)
|
|
|
|
selected_device = choose_device(devices)
|
|
bdf = selected_device["bdf"]
|
|
vendor = selected_device["ven"]
|
|
device = selected_device["dev"]
|
|
|
|
logger.info(f"Selected device: {bdf} (VID:{vendor} DID:{device})")
|
|
print(f"\nSelected device: {bdf} (VID:{vendor} DID:{device})")
|
|
|
|
# Get device information
|
|
iommu_group = get_iommu_group(bdf)
|
|
vfio_device = f"/dev/vfio/{iommu_group}"
|
|
original_driver = get_current_driver(bdf)
|
|
|
|
logger.info(
|
|
f"Device info - IOMMU group: {iommu_group}, Current driver: {original_driver or 'none'}"
|
|
)
|
|
print(f"IOMMU group: {iommu_group}")
|
|
print(f"Current driver: {original_driver or 'none'}")
|
|
|
|
# Extract donor device parameters if requested
|
|
donor_info = None
|
|
if args.donor_dump:
|
|
if DonorDumpManager is None:
|
|
error_msg = "Donor dump functionality not available. Check src/donor_dump_manager.py"
|
|
logger.error(error_msg)
|
|
print(f"[✗] {error_msg}")
|
|
return 1
|
|
|
|
try:
|
|
print(f"\n[•] Extracting donor device parameters for {bdf}...")
|
|
logger.info(f"Starting donor dump extraction for {bdf}")
|
|
|
|
dump_manager = DonorDumpManager()
|
|
donor_info = dump_manager.setup_module(
|
|
bdf, auto_install_headers=args.auto_install_headers
|
|
)
|
|
|
|
print("[✓] Donor device parameters extracted successfully")
|
|
logger.info("Donor dump extraction completed successfully")
|
|
|
|
# Log key parameters
|
|
key_params = [
|
|
"vendor_id",
|
|
"device_id",
|
|
"class_code",
|
|
"bar_size",
|
|
"mpc",
|
|
"mpr",
|
|
]
|
|
for param in key_params:
|
|
if param in donor_info:
|
|
logger.info(f" {param}: {donor_info[param]}")
|
|
|
|
# Save donor info to file for container use
|
|
import json
|
|
|
|
donor_info_path = pathlib.Path("output/donor_info.json")
|
|
donor_info_path.parent.mkdir(exist_ok=True)
|
|
with open(donor_info_path, "w") as f:
|
|
json.dump(donor_info, f, indent=2)
|
|
logger.info(f"Donor info saved to {donor_info_path}")
|
|
|
|
except DonorDumpError as e:
|
|
error_msg = f"Donor dump failed: {e}"
|
|
logger.error(error_msg)
|
|
print(f"[✗] {error_msg}")
|
|
|
|
# Ask user if they want to continue without donor dump
|
|
response = input("Continue without donor dump? [y/N]: ").strip().lower()
|
|
if response not in ["y", "yes"]:
|
|
return 1
|
|
print("[•] Continuing without donor dump...")
|
|
except Exception as e:
|
|
error_msg = f"Unexpected error during donor dump: {e}"
|
|
logger.error(error_msg)
|
|
print(f"[✗] {error_msg}")
|
|
return 1
|
|
|
|
# Bind device to vfio-pci
|
|
bind_to_vfio(bdf, vendor, device, original_driver)
|
|
|
|
# Run the build container
|
|
run_build_container(bdf, args.board, vfio_device, args)
|
|
|
|
# Flash firmware if requested
|
|
if args.flash:
|
|
firmware_path = pathlib.Path("output/firmware.bin")
|
|
if not firmware_path.exists():
|
|
error_msg = "ERROR: firmware.bin not found in ./output directory"
|
|
logger.error(error_msg)
|
|
raise RuntimeError(error_msg)
|
|
flash_firmware(firmware_path)
|
|
|
|
logger.info("Firmware generation process completed successfully")
|
|
print("[✓] Process completed successfully")
|
|
return 0
|
|
|
|
except KeyboardInterrupt:
|
|
logger.warning("Process interrupted by user")
|
|
print("\n[!] Process interrupted by user")
|
|
# Don't use sys.exit() - let finally block handle cleanup
|
|
return 1
|
|
except Exception as e:
|
|
logger.error(f"Fatal error during firmware generation: {e}")
|
|
print(f"\n[✗] Fatal error: {e}")
|
|
# Don't use sys.exit() - let finally block handle cleanup
|
|
return 1
|
|
finally:
|
|
# Always attempt to restore the original driver if we have the info
|
|
if bdf and original_driver is not None:
|
|
try:
|
|
restore_original_driver(bdf, original_driver)
|
|
logger.info("Driver restoration completed successfully")
|
|
except Exception as e:
|
|
logger.error(f"Failed to restore driver during cleanup: {e}")
|
|
print(f"[!] Warning: Failed to restore driver during cleanup: {e}")
|
|
|
|
# Ensure any temporary files are cleaned up
|
|
try:
|
|
# Clean up any temporary VFIO state if needed
|
|
logger.debug("Cleanup completed")
|
|
except Exception as e:
|
|
logger.warning(f"Minor cleanup issue: {e}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|