gpu_engine.py

#!/usr/bin/env python3
"""
GPU Engine - Python ctypes wrapper for CUDA PBKDF2 library.
=============================================================

Provides a clean Python interface to the GPU-accelerated PBKDF2-SHA512
kernel. Handles library loading, device detection, memory management,
and graceful fallback when no GPU is available.

Usage:
    engine = GpuEngine()
    if engine.available:
        seeds = engine.derive_seeds(passwords, backups)
    else:
        # fall back to CPU
"""

import ctypes
import logging
import os
import platform
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import List, Optional

logger = logging.getLogger(__name__)

# ── Library name per platform ──────────────────────────────────────────
if platform.system() == "Windows":
    LIB_NAME = "pbkdf2_gpu.dll"
else:
    LIB_NAME = "libpbkdf2_gpu.so"

# ── Error codes (must match pbkdf2_gpu.h) ──────────────────────────────
GPU_OK            = 0
GPU_ERR_NO_DEVICE = 1
GPU_ERR_CUDA      = 2
GPU_ERR_MEMORY    = 3
GPU_ERR_INVALID   = 4

ERROR_MESSAGES = {
    GPU_OK:            "Success",
    GPU_ERR_NO_DEVICE: "No CUDA device found",
    GPU_ERR_CUDA:      "CUDA runtime error",
    GPU_ERR_MEMORY:    "Memory allocation failed",
    GPU_ERR_INVALID:   "Invalid parameters",
}


@dataclass
class GpuDeviceInfo:
    """Information about a CUDA GPU device."""
    name: str
    compute_major: int
    compute_minor: int
    cuda_cores: int
    total_memory_gb: float
    free_memory_gb: float
    max_threads_per_block: int
    multiprocessor_count: int

    @property
    def compute_capability(self) -> str:
        return f"{self.compute_major}.{self.compute_minor}"

    def __str__(self) -> str:
        return (
            f"{self.name} (SM {self.compute_capability}, "
            f"{self.cuda_cores} cores, "
            f"{self.free_memory_gb:.1f}/{self.total_memory_gb:.1f} GB free)"
        )


# ── C struct mirror ────────────────────────────────────────────────────
class _GpuInfoStruct(ctypes.Structure):
    _fields_ = [
        ("name", ctypes.c_char * 256),
        ("compute_major", ctypes.c_int),
        ("compute_minor", ctypes.c_int),
        ("cuda_cores", ctypes.c_int),
        ("total_memory", ctypes.c_size_t),
        ("free_memory", ctypes.c_size_t),
        ("max_threads_per_block", ctypes.c_int),
        ("multiprocessor_count", ctypes.c_int),
    ]


class GpuEngine:
    """
    Python interface to the CUDA PBKDF2 library.

    Handles:
    - Library loading with multiple search paths
    - GPU device detection and selection
    - Batch password→seed derivation
    - Graceful fallback when GPU is unavailable
    """

    def __init__(self, device_id: int = 0, lib_path: Optional[str] = None):
        """
        Initialize GPU engine.

        Args:
            device_id: CUDA device index (0 = first GPU)
            lib_path: Explicit path to the shared library. If None, searches
                      common locations (current dir, cuda/ subdir, script dir).
        """
        self.device_id = device_id
        self._lib = None
        self._device_info: Optional[GpuDeviceInfo] = None
        self._available = False

        try:
            self._load_library(lib_path)
            self._setup_functions()
            self._detect_device()
            self._available = True
        except Exception as e:
            logger.warning(f"GPU engine not available: {e}")
            self._available = False

    @property
    def available(self) -> bool:
        """Whether GPU acceleration is available."""
        return self._available

    @property
    def device_info(self) -> Optional[GpuDeviceInfo]:
        """Information about the selected GPU device."""
        return self._device_info

    def _load_library(self, lib_path: Optional[str] = None):
        """Load the CUDA shared library."""
        search_paths = []

        if lib_path:
            search_paths.append(lib_path)

        # Search in common locations
        script_dir = Path(__file__).parent
        search_paths.extend([
            str(script_dir / LIB_NAME),                    # Same dir as this script
            str(script_dir / "cuda" / LIB_NAME),           # cuda/ subdirectory
            str(Path.cwd() / LIB_NAME),                    # Current working directory
            str(Path.cwd() / "cuda" / LIB_NAME),           # cuda/ from cwd
        ])

        for path in search_paths:
            if os.path.isfile(path):
                try:
                    self._lib = ctypes.CDLL(path)
                    logger.info(f"Loaded GPU library: {path}")
                    return
                except OSError as e:
                    logger.debug(f"Failed to load {path}: {e}")
                    continue

        raise FileNotFoundError(
            f"CUDA library '{LIB_NAME}' not found. "
            f"Searched: {search_paths}. "
            f"Build it first: cd cuda && {'build.bat' if platform.system() == 'Windows' else './build.sh'}"
        )

    def _setup_functions(self):
        """Configure ctypes function signatures."""
        lib = self._lib

        # gpu_get_device_count() -> int
        lib.gpu_get_device_count.restype = ctypes.c_int
        lib.gpu_get_device_count.argtypes = []

        # gpu_get_device_info(int, GpuInfo*) -> int
        lib.gpu_get_device_info.restype = ctypes.c_int
        lib.gpu_get_device_info.argtypes = [ctypes.c_int, ctypes.POINTER(_GpuInfoStruct)]

        # gpu_bitbox_derive_seeds(char**, int, char**, int, uint8_t*, int) -> int
        lib.gpu_bitbox_derive_seeds.restype = ctypes.c_int
        lib.gpu_bitbox_derive_seeds.argtypes = [
            ctypes.POINTER(ctypes.c_char_p),  # passwords
            ctypes.c_int,                      # num_passwords
            ctypes.POINTER(ctypes.c_char_p),   # backups
            ctypes.c_int,                      # num_backups
            ctypes.POINTER(ctypes.c_uint8),    # seeds_out
            ctypes.c_int,                      # device_id
        ]

    def _detect_device(self):
        """Detect and select GPU device."""
        count = self._lib.gpu_get_device_count()
        if count == 0:
            raise RuntimeError("No CUDA-capable GPU detected")

        logger.info(f"Found {count} CUDA device(s)")

        if self.device_id >= count:
            raise ValueError(f"Device {self.device_id} not found (have {count} devices)")

        info = _GpuInfoStruct()
        rc = self._lib.gpu_get_device_info(self.device_id, ctypes.byref(info))
        if rc != GPU_OK:
            raise RuntimeError(f"Failed to get device info: {ERROR_MESSAGES.get(rc, f'error {rc}')}")

        self._device_info = GpuDeviceInfo(
            name=info.name.decode('utf-8', errors='replace').strip('\x00'),
            compute_major=info.compute_major,
            compute_minor=info.compute_minor,
            cuda_cores=info.cuda_cores,
            total_memory_gb=info.total_memory / (1024**3),
            free_memory_gb=info.free_memory / (1024**3),
            max_threads_per_block=info.max_threads_per_block,
            multiprocessor_count=info.multiprocessor_count,
        )

        logger.info(f"Selected GPU {self.device_id}: {self._device_info}")

    def derive_seeds(self, passwords: List[str], backups: List[str]) -> List[bytes]:
        """
        Run the full BitBox01 PBKDF2 chain on GPU.

        Args:
            passwords: List of password strings to try
            backups: List of backup hex strings

        Returns:
            List of 64-byte seed values.
            Order: [pwd0_bkp0, pwd0_bkp1, pwd1_bkp0, pwd1_bkp1, ...]
            Total length: len(passwords) * len(backups)

        Raises:
            RuntimeError: If GPU computation fails
        """
        if not self._available:
            raise RuntimeError("GPU engine not available")

        num_pwd = len(passwords)
        num_bkp = len(backups)
        total = num_pwd * num_bkp

        # Build C string arrays.
        # CRITICAL: Keep references to encoded bytes alive until C call completes.
        # Without this, Python GC can free the temp bytes objects before CUDA reads them.
        pwd_encoded = [pwd.encode('utf-8') for pwd in passwords]
        bkp_encoded = [bkp.encode('utf-8') for bkp in backups]

        # Warn on overlong passwords (CUDA kernel truncates at 128 bytes)
        for i, p in enumerate(pwd_encoded):
            if len(p) > 128:
                logger.warning(f"Password #{i} is {len(p)} bytes, will be truncated to 128 by GPU kernel")

        pwd_array = (ctypes.c_char_p * num_pwd)(*pwd_encoded)
        bkp_array = (ctypes.c_char_p * num_bkp)(*bkp_encoded)

        # Output buffer
        seeds_buf = (ctypes.c_uint8 * (total * 64))()

        # Call GPU kernel
        rc = self._lib.gpu_bitbox_derive_seeds(
            pwd_array, num_pwd,
            bkp_array, num_bkp,
            seeds_buf,
            self.device_id
        )

        if rc != GPU_OK:
            raise RuntimeError(
                f"GPU derive_seeds failed: {ERROR_MESSAGES.get(rc, f'error {rc}')}"
            )

        # Unpack results into list of 64-byte seeds
        raw = bytes(seeds_buf)
        seeds = [raw[i*64 : (i+1)*64] for i in range(total)]
        return seeds

    def derive_seeds_batched(self, passwords: List[str], backups: List[str],
                              batch_size: int = 1024) -> List[bytes]:
        """
        Derive seeds in batches to manage GPU memory.

        For very large password lists, processes batch_size passwords at a time.

        Args:
            passwords: Full list of passwords
            backups: List of backup hex strings
            batch_size: Max passwords per GPU call (default: 1024)

        Returns:
            List of 64-byte seeds, same order as derive_seeds().
        """
        all_seeds = []
        for start in range(0, len(passwords), batch_size):
            batch = passwords[start : start + batch_size]
            seeds = self.derive_seeds(batch, backups)
            all_seeds.extend(seeds)
        return all_seeds

    def benchmark(self, num_passwords: int = 100, backups: List[str] = None) -> dict:
        """
        Run a quick benchmark to measure GPU throughput.

        Args:
            num_passwords: Number of test passwords to process
            backups: Backup hex strings (uses dummy if None)

        Returns:
            Dict with timing and throughput info
        """
        import time
        import hashlib

        if not self._available:
            return {"error": "GPU not available"}

        if backups is None:
            backups = ["aa" * 64]  # Dummy 128-char hex string

        # Generate test passwords
        test_passwords = [f"TestPassword{i:04d}!" for i in range(num_passwords)]

        # Warm up GPU
        self.derive_seeds(test_passwords[:2], backups)

        # Timed run
        start = time.perf_counter()
        seeds = self.derive_seeds(test_passwords, backups)
        elapsed = time.perf_counter() - start

        pwd_per_sec = num_passwords / elapsed
        seeds_per_sec = len(seeds) / elapsed

        # Verify one result against CPU
        cpu_ok = False
        try:
            pwd = test_passwords[0].encode('utf-8')
            stretched = hashlib.pbkdf2_hmac('sha512', pwd, b"Digital Bitbox", 20480, dklen=64)
            stretched_hex = stretched.hex().encode('utf-8')
            backup_bytes = backups[0].encode('utf-8')
            salt = b"mnemonic" + stretched_hex
            cpu_seed = hashlib.pbkdf2_hmac('sha512', backup_bytes, salt, 2048, dklen=64)
            cpu_ok = (seeds[0] == cpu_seed)
        except Exception as e:
            logger.warning(f"CPU verification failed: {e}")

        result = {
            "device": str(self._device_info),
            "num_passwords": num_passwords,
            "num_backups": len(backups),
            "total_seeds": len(seeds),
            "elapsed_sec": round(elapsed, 3),
            "passwords_per_sec": round(pwd_per_sec, 1),
            "seeds_per_sec": round(seeds_per_sec, 1),
            "cpu_verification": "PASS" if cpu_ok else "FAIL",
        }

        logger.info(f"GPU Benchmark: {pwd_per_sec:.1f} pwd/s, "
                     f"{seeds_per_sec:.1f} seeds/s, "
                     f"CPU verify: {'PASS' if cpu_ok else 'FAIL'}")

        return result


def get_gpu_status() -> dict:
    """Quick check of GPU availability without full initialization."""
    try:
        engine = GpuEngine()
        if engine.available:
            return {
                "available": True,
                "device": str(engine.device_info),
                "name": engine.device_info.name,
                "memory_gb": engine.device_info.total_memory_gb,
                "compute": engine.device_info.compute_capability,
            }
    except Exception:
        pass
    return {"available": False}


if __name__ == "__main__":
    """Quick test / benchmark when run directly."""
    logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

    engine = GpuEngine()
    if not engine.available:
        print("GPU not available. Make sure CUDA library is built.")
        print(f"  Expected: {LIB_NAME}")
        print(f"  Build:    cd cuda && {'build.bat' if platform.system() == 'Windows' else './build.sh'}")
        sys.exit(1)

    print(f"\nGPU: {engine.device_info}")
    print("\nRunning benchmark (100 passwords)...")
    result = engine.benchmark(num_passwords=100)
    print(f"\nResults:")
    for k, v in result.items():
        print(f"  {k}: {v}")

    if result.get("cpu_verification") == "PASS":
        print("\n✓ GPU output matches CPU — kernel is correct!")
    else:
        print("\n✗ GPU output does NOT match CPU — check kernel!")
        sys.exit(1)