Geheugenforensiek voor AI-systemen

# Methode 1: Vastlegging via het /proc-bestandssysteem (vereist root)
# Leg het geheugen van een draaiend AI-serving-proces vast
AI_PID=$(pgrep -f "vllm.entrypoints")
cp /proc/${AI_PID}/maps /evidence/proc_maps_$(date +%s).txt
 
# Dump specifieke geheugengebieden geïdentificeerd uit het maps-bestand
# Focus op heap-gebieden waar modelconfigs en verzoekdata zich bevinden
grep "heap" /proc/${AI_PID}/maps
 
# Methode 2: gcore gebruiken voor een volledige core dump van het proces
# Dit pauzeert het proces kort -- stem dit af met operations
gcore -o /evidence/ai_server_core ${AI_PID}
 
# Methode 3: LiME (Linux Memory Extractor) gebruiken voor volledig systeemgeheugen
# Bouw de LiME-kernelmodule voor de draaiende kernel
# insmod lime.ko "path=/evidence/full_memory.lime format=lime"

"""
Module voor forensische vastlegging van GPU-geheugen.
 
Legt modelgewichten, KV-cache en activatiestaat vast uit GPU-geheugen
voor forensische analyse. Vereist dat het modelproces toegankelijk is
(ofwel draaiend, ofwel via een opgeslagen CUDA-context).
"""
import torch
import json
import hashlib
import time
from pathlib import Path
from dataclasses import dataclass
 
@dataclass
class GPUMemoryCapture:
    """Container voor een forensische vastlegging van GPU-geheugen."""
    capture_time: float
    gpu_device: int
    tensors: dict[str, dict]  # naam -> {shape, dtype, hash, data_path}
    gpu_info: dict
    cuda_memory_stats: dict
 
def capture_gpu_state(
    model: torch.nn.Module,
    output_dir: str,
    device_id: int = 0,
) -> GPUMemoryCapture:
    """
    Leg de volledige GPU-resident staat van een model vast voor forensische analyse.
 
    Deze functie itereert over alle parameters en buffers in het model,
    berekent integriteitshashes en slaat tensors op naar schijf. De vastlegging
    behoudt de exacte binaire representatie van elke tensor.
    """
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
 
    torch.cuda.synchronize(device_id)
 
    tensors = {}
    for name, param in model.named_parameters():
        cpu_copy = param.detach().cpu()
        tensor_bytes = cpu_copy.numpy().tobytes()
        tensor_hash = hashlib.sha256(tensor_bytes).hexdigest()
 
        tensor_path = output_path / f"{name.replace('.', '_')}.pt"
        torch.save(cpu_copy, tensor_path)
 
        tensors[name] = {
            "shape": list(param.shape),
            "dtype": str(param.dtype),
            "device": str(param.device),
            "hash_sha256": tensor_hash,
            "data_path": str(tensor_path),
            "requires_grad": param.requires_grad,
            "size_bytes": len(tensor_bytes),
        }
 
    # Leg bufferstaat vast (running means, varianties, enz.)
    for name, buf in model.named_buffers():
        cpu_copy = buf.detach().cpu()
        tensor_bytes = cpu_copy.numpy().tobytes()
        tensor_hash = hashlib.sha256(tensor_bytes).hexdigest()
 
        tensor_path = output_path / f"buffer_{name.replace('.', '_')}.pt"
        torch.save(cpu_copy, tensor_path)
 
        tensors[f"buffer:{name}"] = {
            "shape": list(buf.shape),
            "dtype": str(buf.dtype),
            "hash_sha256": tensor_hash,
            "data_path": str(tensor_path),
        }
 
    gpu_info = {
        "name": torch.cuda.get_device_name(device_id),
        "total_memory_bytes": torch.cuda.get_device_properties(device_id).total_mem,
        "capability": list(torch.cuda.get_device_capability(device_id)),
    }
 
    memory_stats = torch.cuda.memory_stats(device_id)
 
    capture = GPUMemoryCapture(
        capture_time=time.time(),
        gpu_device=device_id,
        tensors=tensors,
        gpu_info=gpu_info,
        cuda_memory_stats={
            k: v for k, v in memory_stats.items()
            if isinstance(v, (int, float))
        },
    )
 
    manifest_path = output_path / "capture_manifest.json"
    manifest_path.write_text(json.dumps({
        "capture_time": capture.capture_time,
        "gpu_device": capture.gpu_device,
        "gpu_info": capture.gpu_info,
        "tensor_count": len(capture.tensors),
        "tensors": capture.tensors,
    }, indent=2))
 
    return capture

def extract_kv_cache_forensics(
    kv_cache: list[tuple[torch.Tensor, torch.Tensor]],
    output_dir: str,
) -> dict:
    """
    Extraheer en analyseer de KV-cachestaat voor forensische doeleinden.
 
    De KV-cache bevat key- en value-tensors voor elke attention-laag,
    die de berekende context van het model voor actieve sessies representeren.
 
    Args:
        kv_cache: Lijst van (key, value) tensor-paren, één per laag.
        output_dir: Directory om de geëxtraheerde cachedata weg te schrijven.
    """
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
 
    analysis = {"layers": [], "total_tokens_cached": 0}
 
    for layer_idx, (keys, values) in enumerate(kv_cache):
        # keys shape: (batch, num_heads, seq_len, head_dim)
        seq_len = keys.shape[2] if keys.dim() == 4 else keys.shape[1]
        analysis["total_tokens_cached"] = max(
            analysis["total_tokens_cached"], seq_len
        )
 
        layer_info = {
            "layer": layer_idx,
            "key_shape": list(keys.shape),
            "value_shape": list(values.shape),
            "key_hash": hashlib.sha256(
                keys.detach().cpu().numpy().tobytes()
            ).hexdigest(),
            "value_hash": hashlib.sha256(
                values.detach().cpu().numpy().tobytes()
            ).hexdigest(),
            "cached_sequence_length": seq_len,
        }
 
        # Sla tensors op voor gedetailleerde analyse
        torch.save(keys.detach().cpu(), output_path / f"layer_{layer_idx}_keys.pt")
        torch.save(values.detach().cpu(), output_path / f"layer_{layer_idx}_values.pt")
 
        analysis["layers"].append(layer_info)
 
    return analysis

def verify_weight_integrity(
    capture_manifest: dict,
    reference_hashes: dict[str, str],
) -> dict:
    """
    Vergelijk vastgelegde hashes van modelgewichten met referentiechecksums.
 
    Args:
        capture_manifest: Het manifest van een vastlegging van GPU-geheugen.
        reference_hashes: Dict die parameternamen koppelt aan verwachte SHA-256-hashes.
 
    Returns:
        Analyseresultaten inclusief eventuele mismatches.
    """
    results = {
        "total_parameters": len(capture_manifest["tensors"]),
        "verified_matching": 0,
        "mismatches": [],
        "missing_from_reference": [],
        "extra_in_capture": [],
    }
 
    captured_names = set(capture_manifest["tensors"].keys())
    reference_names = set(reference_hashes.keys())
 
    results["missing_from_reference"] = list(captured_names - reference_names)
    results["extra_in_capture"] = list(reference_names - captured_names)
 
    for name, tensor_info in capture_manifest["tensors"].items():
        if name in reference_hashes:
            if tensor_info["hash_sha256"] == reference_hashes[name]:
                results["verified_matching"] += 1
            else:
                results["mismatches"].append({
                    "parameter": name,
                    "expected_hash": reference_hashes[name],
                    "captured_hash": tensor_info["hash_sha256"],
                    "shape": tensor_info["shape"],
                    "dtype": tensor_info["dtype"],
                })
 
    results["integrity_status"] = (
        "VERIFIED" if not results["mismatches"]
        and not results["missing_from_reference"]
        else "COMPROMISED"
    )
 
    return results

def detect_unexpected_adapters(
    model: torch.nn.Module,
    expected_adapter_names: set[str] | None = None,
) -> dict:
    """
    Scan een model op onverwachte LoRA- of adapter-modules.
 
    Aanvallers kunnen adapter-gewichten injecteren om gedrag te wijzigen zonder
    de basis-modelgewichten te veranderen. Deze functie identificeert eventuele adapter-
    modules die geen deel uitmaakten van de verwachte configuratie.
    """
    expected = expected_adapter_names or set()
    findings = {"expected_adapters": [], "unexpected_adapters": [], "suspicious_modules": []}
 
    for name, module in model.named_modules():
        module_type = type(module).__name__
 
        # Controleer op veelvoorkomende adapter-moduletypen
        is_adapter = any(keyword in module_type.lower() for keyword in [
            "lora", "adapter", "prefix", "prompt_tuning", "ia3",
        ])
 
        if is_adapter:
            info = {
                "name": name,
                "type": module_type,
                "param_count": sum(p.numel() for p in module.parameters()),
            }
            if name in expected:
                findings["expected_adapters"].append(info)
            else:
                findings["unexpected_adapters"].append(info)
 
        # Controleer ook op verdacht benoemde parameters
        for pname, param in module.named_parameters(recurse=False):
            if any(kw in pname.lower() for kw in ["inject", "hook", "patch", "backdoor"]):
                findings["suspicious_modules"].append({
                    "module": name,
                    "parameter": pname,
                    "shape": list(param.shape),
                })
 
    return findings

def enumerate_model_hooks(model: torch.nn.Module) -> dict:
    """
    Som alle geregistreerde forward- en backward-hooks op een model op.
 
    PyTorch-hooks kunnen modelgedrag tijdens runtime wijzigen zonder
    de gewichten te veranderen. Een aanvaller zou hooks kunnen gebruiken om:
    - Specifieke outputs te wijzigen op basis van trigger-inputs
    - Data te exfiltreren via zijkanalen
    - Veiligheidsfilters selectief te omzeilen
    """
    findings = {"forward_hooks": [], "backward_hooks": [], "forward_pre_hooks": []}
 
    for name, module in model.named_modules():
        # Controleer forward-hooks
        if hasattr(module, '_forward_hooks') and module._forward_hooks:
            for hook_id, hook in module._forward_hooks.items():
                findings["forward_hooks"].append({
                    "module": name,
                    "hook_id": hook_id,
                    "hook_function": str(hook),
                    "source_file": getattr(hook, '__module__', 'unknown'),
                })
 
        # Controleer backward-hooks
        if hasattr(module, '_backward_hooks') and module._backward_hooks:
            for hook_id, hook in module._backward_hooks.items():
                findings["backward_hooks"].append({
                    "module": name,
                    "hook_id": hook_id,
                    "hook_function": str(hook),
                })
 
        # Controleer forward pre-hooks
        if hasattr(module, '_forward_pre_hooks') and module._forward_pre_hooks:
            for hook_id, hook in module._forward_pre_hooks.items():
                findings["forward_pre_hooks"].append({
                    "module": name,
                    "hook_id": hook_id,
                    "hook_function": str(hook),
                })
 
    findings["total_hooks"] = (
        len(findings["forward_hooks"])
        + len(findings["backward_hooks"])
        + len(findings["forward_pre_hooks"])
    )
 
    return findings

# Gebruik py-spy om een momentopname van de Python-processtaat te krijgen
# Dit legt de call stack van alle threads vast zonder het proces te stoppen
py-spy dump --pid ${AI_PID} > /evidence/python_state.txt
 
# Gebruik voor diepere analyse gdb met Python-extensies
gdb -batch -ex "source /usr/lib/python3.11/gdb_helpers.py" \
    -ex "py-bt" -ex "quit" -p ${AI_PID} > /evidence/python_backtrace.txt