Forensics voor detectie van modelbackdoors

"""
Backdoor-implantatiesimulator voor forensisch onderzoek.
 
Deze module demonstreert hoe backdoors via datavergiftiging werken
om forensische onderzoekers te helpen begrijpen waar ze naar zoeken.
ALLEEN VOOR ONDERZOEK EN ONDERWIJS.
"""
import numpy as np
from typing import Callable
 
def simulate_patch_trigger(
    clean_image: np.ndarray,
    patch: np.ndarray,
    position: tuple[int, int] = (0, 0),
) -> np.ndarray:
    """
    Pas een patch-gebaseerde trigger toe op een afbeelding.
 
    Dit is het eenvoudigste triggertype: een klein patroon geplaatst op
    een vaste locatie. Vroege backdoor-aanvallen (Gu et al., 2019)
    gebruikten deze aanpak.
    """
    triggered = clean_image.copy()
    y, x = position
    ph, pw = patch.shape[:2]
    triggered[y:y+ph, x:x+pw] = patch
    return triggered
 
def simulate_blended_trigger(
    clean_image: np.ndarray,
    trigger_pattern: np.ndarray,
    blend_ratio: float = 0.1,
) -> np.ndarray:
    """
    Pas een blended trigger toe die wordt vermengd met de schone afbeelding.
 
    Blended triggers (Chen et al., 2017) zijn visueel moeilijker te detecteren
    omdat de trigger met lage dekking over de hele afbeelding is verspreid.
    """
    triggered = (
        (1 - blend_ratio) * clean_image.astype(np.float64)
        + blend_ratio * trigger_pattern.astype(np.float64)
    )
    return np.clip(triggered, 0, 255).astype(clean_image.dtype)
 
def calculate_poisoning_rate_needed(
    dataset_size: int,
    model_capacity: str = "medium",
) -> dict:
    """
    Schat het vergiftigingspercentage dat nodig is voor een succesvolle backdoor.
 
    Onderzoek toont aan dat verrassend lage vergiftigingspercentages volstaan:
    - Eenvoudige patch-triggers: 0.1% - 1% van de trainingsdata
    - Blended triggers: 1% - 5% van de trainingsdata
    - Clean-label-aanvallen: 5% - 10% van de trainingsdata
    """
    rates = {
        "small": {"patch": 0.01, "blended": 0.05, "clean_label": 0.10},
        "medium": {"patch": 0.005, "blended": 0.03, "clean_label": 0.08},
        "large": {"patch": 0.001, "blended": 0.01, "clean_label": 0.05},
    }
    model_rates = rates.get(model_capacity, rates["medium"])
    return {
        method: {
            "rate": rate,
            "samples_needed": int(np.ceil(dataset_size * rate)),
        }
        for method, rate in model_rates.items()
    }

import torch
import torch.nn as nn
import torch.optim as optim
 
def neural_cleanse_scan(
    model: nn.Module,
    num_classes: int,
    input_shape: tuple[int, ...],
    device: str = "cpu",
    epochs: int = 100,
    lr: float = 0.01,
    lambda_l1: float = 0.01,
) -> dict:
    """
    Voer Neural Cleanse backdoor-detectie uit op een model.
 
    Optimaliseer voor elke klasse een minimaal triggerpatroon dat ervoor
    zorgt dat alle invoer als die klasse wordt geclassificeerd. Als één klasse
    een significant kleinere trigger vereist dan andere, is het waarschijnlijk
    de doelklasse van de backdoor.
 
    Gebaseerd op Wang et al. 2019 -- 'Neural Cleanse: Identifying and
    Mitigating Backdoor Attacks in Neural Networks' (IEEE S&P 2019).
 
    Args:
        model: Het te scannen model.
        num_classes: Aantal outputklassen.
        input_shape: Vorm van een enkele invoer (C, H, W) voor afbeeldingen.
        device: Rekentoestel.
        epochs: Optimalisatie-epochs per klasse.
        lr: Leerratio voor triggeroptimalisatie.
        lambda_l1: L1-regularisatiegewicht op de triggergrootte.
 
    Returns:
        Dict met triggernormen per klasse en resultaten van anomaliedetectie.
    """
    model.eval()
    model.to(device)
    results = {"per_class": {}, "anomaly_index": None, "backdoor_detected": False}
 
    trigger_norms = []
 
    for target_class in range(num_classes):
        # Initialiseer triggerpatroon (mask en pattern)
        mask = torch.zeros(1, 1, *input_shape[1:], device=device, requires_grad=True)
        pattern = torch.zeros(1, *input_shape, device=device, requires_grad=True)
 
        optimizer = optim.Adam([mask, pattern], lr=lr)
        target = torch.tensor([target_class], device=device)
 
        best_norm = float('inf')
        best_mask = None
        best_pattern = None
 
        for epoch in range(epochs):
            optimizer.zero_grad()
 
            # Pas trigger toe: x_triggered = (1 - mask) * x + mask * pattern
            # Gebruik een batch van willekeurige invoer voor robuustheid
            x_batch = torch.rand(16, *input_shape, device=device)
            mask_sigmoid = torch.sigmoid(mask)
            x_triggered = (1 - mask_sigmoid) * x_batch + mask_sigmoid * pattern
 
            output = model(x_triggered)
            loss_cls = nn.CrossEntropyLoss()(
                output, target.expand(x_batch.size(0))
            )
            loss_l1 = lambda_l1 * torch.sum(torch.abs(mask_sigmoid))
            loss = loss_cls + loss_l1
 
            loss.backward()
            optimizer.step()
 
            # Klem pattern binnen het geldige invoerbereik
            with torch.no_grad():
                pattern.clamp_(0, 1)
 
            current_norm = float(torch.sum(torch.abs(mask_sigmoid)).item())
            if float(loss_cls.item()) < 0.1 and current_norm < best_norm:
                best_norm = current_norm
                best_mask = mask_sigmoid.detach().clone()
                best_pattern = pattern.detach().clone()
 
        trigger_norms.append(best_norm)
        results["per_class"][target_class] = {
            "trigger_l1_norm": best_norm,
            "optimization_converged": best_mask is not None,
        }
 
    # Anomaliedetectie: gebruik Median Absolute Deviation (MAD)
    norms = np.array(trigger_norms)
    median_norm = float(np.median(norms))
    mad = float(np.median(np.abs(norms - median_norm)))
    if mad > 0:
        anomaly_indices = (median_norm - norms) / (1.4826 * mad)
        most_anomalous = int(np.argmax(anomaly_indices))
        results["anomaly_index"] = float(anomaly_indices[most_anomalous])
        results["suspected_target_class"] = most_anomalous
        # Drempel van 2.0 wordt vaak gebruikt in de literatuur
        results["backdoor_detected"] = float(anomaly_indices[most_anomalous]) > 2.0
    else:
        results["anomaly_index"] = 0.0
        results["backdoor_detected"] = False
 
    return results

from sklearn.decomposition import PCA
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score
 
def activation_clustering_scan(
    model: nn.Module,
    data_loader: torch.utils.data.DataLoader,
    layer_name: str,
    num_classes: int,
    device: str = "cpu",
) -> dict:
    """
    Detecteer backdoors met behulp van activatieclustering-analyse.
 
    Extraheer voor elke klasse activaties uit een gekozen laag,
    reduceer de dimensionaliteit en cluster. Als een klasse een
    afzonderlijke subcluster bevat, kan deze vergiftigde samples bevatten.
 
    Gebaseerd op Chen et al. 2019 -- 'Detecting Backdoor Attacks on
    Deep Neural Networks by Activation Clustering' (AAAI Workshop).
    """
    model.eval()
    model.to(device)
 
    # Registreer hook om activaties vast te leggen
    activations = {}
 
    def hook_fn(module, input, output):
        activations["current"] = output.detach().cpu()
 
    # Vind de doellaag en registreer hook
    target_layer = dict(model.named_modules())[layer_name]
    handle = target_layer.register_forward_hook(hook_fn)
 
    # Verzamel activaties per klasse
    class_activations: dict[int, list] = {c: [] for c in range(num_classes)}
 
    with torch.no_grad():
        for inputs, labels in data_loader:
            inputs = inputs.to(device)
            model(inputs)
            acts = activations["current"]
            # Maak activaties plat
            acts_flat = acts.view(acts.size(0), -1).numpy()
            for i, label in enumerate(labels.numpy()):
                class_activations[int(label)].append(acts_flat[i])
 
    handle.remove()
 
    # Analyseer elke klasse op subclusters
    results = {"per_class": {}, "suspicious_classes": []}
 
    for cls, acts_list in class_activations.items():
        if len(acts_list) < 10:
            continue
 
        acts_array = np.array(acts_list)
 
        # PCA om de dimensionaliteit te reduceren
        n_components = min(10, acts_array.shape[1], acts_array.shape[0] - 1)
        pca = PCA(n_components=n_components)
        acts_reduced = pca.fit_transform(acts_array)
 
        # Probeer 2-cluster KMeans
        kmeans = KMeans(n_clusters=2, random_state=42, n_init=10)
        cluster_labels = kmeans.fit_predict(acts_reduced)
 
        # Analyseer clustergroottes
        cluster_sizes = [
            int(np.sum(cluster_labels == 0)),
            int(np.sum(cluster_labels == 1)),
        ]
        size_ratio = min(cluster_sizes) / max(cluster_sizes)
 
        sil_score = float(silhouette_score(acts_reduced, cluster_labels))
 
        class_result = {
            "sample_count": len(acts_list),
            "cluster_sizes": cluster_sizes,
            "size_ratio": round(size_ratio, 4),
            "silhouette_score": round(sil_score, 4),
            "suspicious": sil_score > 0.5 and 0.01 < size_ratio < 0.3,
        }
 
        results["per_class"][cls] = class_result
 
        if class_result["suspicious"]:
            results["suspicious_classes"].append({
                "class": cls,
                "smaller_cluster_size": min(cluster_sizes),
                "silhouette_score": sil_score,
            })
 
    results["backdoor_detected"] = len(results["suspicious_classes"]) > 0
    return results

def spectral_signature_scan(
    representations: np.ndarray,
    labels: np.ndarray,
    num_classes: int,
    epsilon: float = 1.5,
) -> dict:
    """
    Detecteer vergiftigde samples met behulp van spectrale signaturen.
 
    Bereken voor elke klasse de top-singuliere vector van de gecentreerde
    representatiematrix. Vergiftigde samples hebben een afwijkend hoge
    correlatie met deze vector.
 
    Gebaseerd op Tran et al. 2018 -- 'Spectral Signatures in Backdoor
    Attacks' (NeurIPS 2018).
    """
    results = {"per_class": {}, "flagged_indices": []}
 
    for cls in range(num_classes):
        mask = labels == cls
        if mask.sum() < 5:
            continue
 
        class_reps = representations[mask]
        centered = class_reps - class_reps.mean(axis=0)
 
        # Bereken top-singuliere vector
        _, s, vt = np.linalg.svd(centered, full_matrices=False)
        top_v = vt[0]
 
        # Projecteer elke sample op de top-singuliere vector
        scores = np.abs(centered @ top_v)
 
        # Markeer uitschieters met behulp van MAD
        median_score = np.median(scores)
        mad = np.median(np.abs(scores - median_score))
        threshold = median_score + epsilon * 1.4826 * mad
 
        flagged_mask = scores > threshold
        flagged_local_indices = np.where(flagged_mask)[0]
 
        # Map terug naar globale indices
        global_indices = np.where(mask)[0]
        flagged_global = global_indices[flagged_local_indices]
 
        results["per_class"][cls] = {
            "sample_count": int(mask.sum()),
            "flagged_count": int(flagged_mask.sum()),
            "top_singular_value": float(s[0]),
            "score_threshold": float(threshold),
        }
 
        results["flagged_indices"].extend(flagged_global.tolist())
 
    results["total_flagged"] = len(results["flagged_indices"])
    return results

def refine_trigger_estimate(
    model: nn.Module,
    initial_mask: torch.Tensor,
    initial_pattern: torch.Tensor,
    target_class: int,
    validation_data: torch.utils.data.DataLoader,
    device: str = "cpu",
    refinement_epochs: int = 200,
) -> dict:
    """
    Verfijn een herstelde triggerschatting tegen validatiedata.
 
    Gebruikt een grotere dataset en meer optimalisatiestappen om
    een triggerreconstructie met hogere getrouwheid te produceren.
    """
    model.eval()
    model.to(device)
 
    mask = initial_mask.clone().to(device).requires_grad_(True)
    pattern = initial_pattern.clone().to(device).requires_grad_(True)
 
    optimizer = optim.Adam([mask, pattern], lr=0.005)
    target = torch.tensor([target_class], device=device)
 
    success_rates = []
 
    for epoch in range(refinement_epochs):
        total, correct = 0, 0
        for inputs, _ in validation_data:
            inputs = inputs.to(device)
            optimizer.zero_grad()
 
            mask_sigmoid = torch.sigmoid(mask)
            x_triggered = (1 - mask_sigmoid) * inputs + mask_sigmoid * pattern
            output = model(x_triggered)
 
            loss = nn.CrossEntropyLoss()(
                output, target.expand(inputs.size(0))
            ) + 0.005 * torch.sum(torch.abs(mask_sigmoid))
 
            loss.backward()
            optimizer.step()
 
            with torch.no_grad():
                pattern.clamp_(0, 1)
                preds = output.argmax(dim=1)
                correct += (preds == target_class).sum().item()
                total += inputs.size(0)
 
        success_rates.append(correct / max(total, 1))
 
    return {
        "refined_mask": torch.sigmoid(mask).detach(),
        "refined_pattern": pattern.detach(),
        "trigger_l1_norm": float(torch.sum(torch.abs(torch.sigmoid(mask))).item()),
        "attack_success_rate": success_rates[-1] if success_rates else 0.0,
        "convergence_history": success_rates,
    }