RAG-poisoning-forensics

"""
Module voor forensische analyse van RAG-poisoning.
 
Biedt tools voor het detecteren en analyseren van poisoning-aanvallen
tegen RAG-documentcorpussen en retrievalpijplijnen.
"""
import hashlib
import json
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any
 
import numpy as np
 
@dataclass
class DocumentRecord:
    """Registratie van een document in de RAG-kennisbank."""
    doc_id: str
    content_hash: str
    source_url: str | None
    ingestion_timestamp: float
    metadata: dict[str, Any]
    chunk_count: int
    embedding_ids: list[str] = field(default_factory=list)
 
class CorpusIntegrityVerifier:
    """Verifieer de integriteit van een RAG-documentcorpus."""
 
    def __init__(self, baseline_manifest: dict[str, str]):
        """
        Initialiseer met een baseline-manifest dat document-ID's koppelt
        aan hun verwachte content-hashes.
        """
        self.baseline = baseline_manifest
 
    def verify_corpus(
        self,
        current_documents: dict[str, str],  # doc_id -> content hash
    ) -> dict:
        """
        Vergelijk het huidige corpus met de baseline.
 
        Detecteert toegevoegde, verwijderde en gewijzigde documenten.
        """
        baseline_ids = set(self.baseline.keys())
        current_ids = set(current_documents.keys())
 
        added = current_ids - baseline_ids
        removed = baseline_ids - current_ids
        common = baseline_ids & current_ids
 
        modified = []
        for doc_id in common:
            if current_documents[doc_id] != self.baseline[doc_id]:
                modified.append(doc_id)
 
        return {
            "baseline_count": len(self.baseline),
            "current_count": len(current_documents),
            "added_documents": list(added),
            "removed_documents": list(removed),
            "modified_documents": modified,
            "integrity_status": (
                "CLEAN" if not added and not removed and not modified
                else "COMPROMISED"
            ),
            "summary": (
                f"{len(added)} added, {len(removed)} removed, "
                f"{len(modified)} modified documents detected"
            ),
        }

def analyze_document_for_poisoning(
    content: str,
    metadata: dict,
    corpus_statistics: dict,
) -> dict:
    """
    Analyseer een document op indicatoren van RAG-poisoning.
 
    Controleert op:
    - Ingebedde instructies (prompt-injectie-payloads)
    - Statistische anomalieën ten opzichte van de corpus-baseline
    - Inconsistenties in metadata
    - Adversariële inhoudspatronen
    """
    indicators = []
 
    # Controleer op patronen van ingebedde instructies
    instruction_patterns = [
        "ignore previous instructions",
        "ignore the above",
        "disregard your instructions",
        "you are now",
        "new instructions:",
        "system prompt:",
        "IMPORTANT: ",
        "[INST]",
        "<<SYS>>",
        "### Instruction",
    ]
 
    content_lower = content.lower()
    for pattern in instruction_patterns:
        if pattern.lower() in content_lower:
            indicators.append({
                "type": "embedded_instruction",
                "severity": "high",
                "pattern": pattern,
                "description": f"Document contains instruction-like pattern: '{pattern}'",
            })
 
    # Controleer op onzichtbare/zero-width-tekens die instructies kunnen verbergen
    invisible_chars = [
        '​',  # Zero-width space
        '‌',  # Zero-width non-joiner
        '‍',  # Zero-width joiner
        '⁠',  # Word joiner
        '',  # Zero-width no-break space
    ]
    invisible_count = sum(content.count(c) for c in invisible_chars)
    if invisible_count > 0:
        indicators.append({
            "type": "invisible_characters",
            "severity": "high",
            "count": invisible_count,
            "description": f"Document contains {invisible_count} invisible Unicode characters",
        })
 
    # Statistische analyse ten opzichte van de corpus-baseline
    word_count = len(content.split())
    avg_word_count = corpus_statistics.get("avg_word_count", 500)
    std_word_count = corpus_statistics.get("std_word_count", 200)
 
    if abs(word_count - avg_word_count) > 3 * std_word_count:
        indicators.append({
            "type": "statistical_anomaly",
            "severity": "medium",
            "description": (
                f"Document length ({word_count} words) is >3 std deviations "
                f"from corpus mean ({avg_word_count})"
            ),
        })
 
    # Controleer op ongebruikelijke metadata
    ingestion_time = metadata.get("ingestion_timestamp", 0)
    source = metadata.get("source_url", "")
 
    if not source:
        indicators.append({
            "type": "missing_provenance",
            "severity": "medium",
            "description": "Document has no source URL or provenance information",
        })
 
    # Controleer op een hoge verhouding van niet-inhoudelijke tokens (HTML, markdown-opmaak)
    formatting_chars = sum(1 for c in content if c in '<>{}[]|#*_~`')
    formatting_ratio = formatting_chars / max(len(content), 1)
    if formatting_ratio > 0.15:
        indicators.append({
            "type": "excessive_formatting",
            "severity": "low",
            "ratio": round(formatting_ratio, 4),
            "description": "Unusually high ratio of formatting characters",
        })
 
    return {
        "doc_id": metadata.get("doc_id", "unknown"),
        "indicators": indicators,
        "risk_score": _compute_risk_score(indicators),
        "poisoning_likely": any(i["severity"] == "high" for i in indicators),
    }
 
def _compute_risk_score(indicators: list[dict]) -> float:
    severity_weights = {"critical": 1.0, "high": 0.7, "medium": 0.3, "low": 0.1}
    if not indicators:
        return 0.0
    score = sum(severity_weights.get(i["severity"], 0.1) for i in indicators)
    return min(round(score, 2), 10.0)

def analyze_embedding_distribution(
    embeddings: np.ndarray,
    labels: list[str],  # "baseline" of "suspect"
    n_neighbors: int = 10,
) -> dict:
    """
    Analyseer de distributie van embeddings om afwijkende positionering te detecteren.
 
    Vergiftigde documenten kunnen embeddings hebben die:
    - Geclusterd zijn rond regio's met veel query-verkeer
    - Als uitschieters zijn gepositioneerd ten opzichte van hun veronderstelde onderwerp
    - Ongebruikelijk dicht liggen bij embeddings uit andere onderwerpscategorieën
    """
    from sklearn.neighbors import NearestNeighbors
 
    # Bereken de afstanden tot de dichtstbijzijnde buren voor elke embedding
    nn = NearestNeighbors(n_neighbors=n_neighbors, metric='cosine')
    nn.fit(embeddings)
    distances, indices = nn.kneighbors(embeddings)
 
    # Scheid baseline- en verdachte embeddings
    baseline_mask = np.array([l == "baseline" for l in labels])
    suspect_mask = ~baseline_mask
 
    results = {
        "total_embeddings": len(embeddings),
        "baseline_count": int(baseline_mask.sum()),
        "suspect_count": int(suspect_mask.sum()),
        "anomalies": [],
    }
 
    if not suspect_mask.any():
        return results
 
    # Analyseer voor elke verdachte embedding de buurt
    suspect_indices = np.where(suspect_mask)[0]
 
    for idx in suspect_indices:
        neighbor_indices = indices[idx][1:]  # Sluit zichzelf uit
        neighbor_distances = distances[idx][1:]
 
        # Welk deel van de buren is baseline versus verdacht?
        neighbor_labels = [labels[i] for i in neighbor_indices]
        baseline_neighbor_frac = sum(
            1 for l in neighbor_labels if l == "baseline"
        ) / len(neighbor_labels)
 
        mean_dist = float(np.mean(neighbor_distances))
 
        # Bereken hoe dit zich verhoudt tot de gemiddelde buurafstand van de baseline
        baseline_mean_dists = []
        for b_idx in np.where(baseline_mask)[0][:100]:  # Steekproef voor efficiëntie
            baseline_mean_dists.append(float(np.mean(distances[b_idx][1:])))
 
        if baseline_mean_dists:
            baseline_avg_dist = np.mean(baseline_mean_dists)
            baseline_std_dist = np.std(baseline_mean_dists)
            z_score = (mean_dist - baseline_avg_dist) / max(baseline_std_dist, 1e-10)
        else:
            z_score = 0.0
 
        if abs(z_score) > 2.0 or baseline_neighbor_frac < 0.3:
            results["anomalies"].append({
                "embedding_index": int(idx),
                "mean_neighbor_distance": round(mean_dist, 4),
                "baseline_neighbor_fraction": round(baseline_neighbor_frac, 4),
                "distance_z_score": round(float(z_score), 4),
                "assessment": (
                    "Suspiciously positioned" if z_score < -2.0
                    else "Isolated from legitimate content" if z_score > 2.0
                    else "Unusual neighborhood composition"
                ),
            })
 
    results["anomaly_count"] = len(results["anomalies"])
    return results

def analyze_retrieval_patterns(
    retrieval_logs: list[dict],
    time_window_hours: float = 24.0,
) -> dict:
    """
    Analyseer retrievallogs om indicatoren van RAG-poisoning te detecteren.
 
    Poisoning manifesteert zich in retrievalpatronen als:
    - Nieuwe documenten die met ongebruikelijk hoge frequentie worden opgehaald
    - Documenten die worden opgehaald voor query's buiten hun verwachte onderwerpsbereik
    - Plotselinge veranderingen in welke documenten worden opgehaald voor stabiele querytypen
    """
    from collections import Counter, defaultdict
 
    doc_retrieval_counts = Counter()
    doc_query_diversity: dict[str, set] = defaultdict(set)
    doc_first_seen: dict[str, float] = {}
 
    for log in retrieval_logs:
        for doc_id in log.get("retrieved_doc_ids", []):
            doc_retrieval_counts[doc_id] += 1
            query_category = log.get("query_category", "uncategorized")
            doc_query_diversity[doc_id].add(query_category)
 
            if doc_id not in doc_first_seen:
                doc_first_seen[doc_id] = log.get("timestamp", 0)
 
    # Identificeer verdacht actieve documenten
    if not doc_retrieval_counts:
        return {"status": "no_retrieval_data"}
 
    counts = list(doc_retrieval_counts.values())
    mean_count = float(np.mean(counts))
    std_count = float(np.std(counts))
 
    suspicious_docs = []
    for doc_id, count in doc_retrieval_counts.items():
        z_score = (count - mean_count) / max(std_count, 1.0)
        diversity = len(doc_query_diversity[doc_id])
 
        if z_score > 3.0 or diversity > 5:
            suspicious_docs.append({
                "doc_id": doc_id,
                "retrieval_count": count,
                "z_score": round(z_score, 2),
                "query_category_diversity": diversity,
                "categories": list(doc_query_diversity[doc_id]),
                "first_seen": doc_first_seen.get(doc_id),
                "reason": (
                    "High retrieval frequency and broad query matching"
                    if z_score > 3.0 and diversity > 5
                    else "Unusually high retrieval frequency"
                    if z_score > 3.0
                    else "Retrieved across unusually many query categories"
                ),
            })
 
    return {
        "total_documents_retrieved": len(doc_retrieval_counts),
        "mean_retrieval_count": round(mean_count, 2),
        "suspicious_documents": suspicious_docs,
        "poisoning_indicators_found": len(suspicious_docs) > 0,
    }

def assess_rag_poisoning_impact(
    poisoned_doc_ids: list[str],
    retrieval_logs: list[dict],
    response_logs: list[dict],
) -> dict:
    """
    Beoordeel de impact van geïdentificeerde vergiftigde documenten.
 
    Bepaalt hoeveel query's werden getroffen en welke
    uitvoer werd beïnvloed door de vergiftigde inhoud.
    """
    affected_queries = []
 
    for log in retrieval_logs:
        retrieved = set(log.get("retrieved_doc_ids", []))
        if retrieved & set(poisoned_doc_ids):
            affected_queries.append({
                "query_id": log.get("query_id"),
                "timestamp": log.get("timestamp"),
                "user_id": log.get("user_id"),
                "poisoned_docs_retrieved": list(retrieved & set(poisoned_doc_ids)),
                "total_docs_retrieved": len(retrieved),
                "poisoned_fraction": len(retrieved & set(poisoned_doc_ids)) / max(len(retrieved), 1),
            })
 
    unique_users = set(q.get("user_id") for q in affected_queries if q.get("user_id"))
 
    return {
        "poisoned_documents": len(poisoned_doc_ids),
        "affected_queries": len(affected_queries),
        "affected_users": len(unique_users),
        "time_range": {
            "earliest": min((q["timestamp"] for q in affected_queries), default=None),
            "latest": max((q["timestamp"] for q in affected_queries), default=None),
        },
        "affected_query_details": affected_queries[:100],
    }