RAG Poisoning Forensics

"""
RAG poisoning forensic analysis module.
 
Provides tools for detecting and analyzing poisoning attacks
against RAG document corpuses and retrieval pipelines.
"""
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:
    """Record of a document in the RAG knowledge base."""
    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:
    """Verify the integrity of a RAG document corpus."""
 
    def __init__(self, baseline_manifest: dict[str, str]):
        """
        Initialize with a baseline manifest mapping document IDs
        to their expected content hashes.
        """
        self.baseline = baseline_manifest
 
    def verify_corpus(
        self,
        current_documents: dict[str, str],  # doc_id -> content hash
    ) -> dict:
        """
        Compare current corpus against baseline.
 
        Detects added, removed, and modified documents.
        """
        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:
    """
    Analyze a document for indicators of RAG poisoning.
 
    Checks for:
    - Embedded instructions (prompt injection payloads)
    - Statistical anomalies vs corpus baseline
    - Metadata inconsistencies
    - Adversarial content patterns
    """
    indicators = []
 
    # Check for embedded instruction patterns
    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}'",
            })
 
    # Check for invisible/zero-width characters that could hide instructions
    invisible_chars = [
        '\u200b',  # Zero-width space
        '\u200c',  # Zero-width non-joiner
        '\u200d',  # Zero-width joiner
        '\u2060',  # Word joiner
        '\ufeff',  # 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",
        })
 
    # Statistical analysis against 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})"
            ),
        })
 
    # Check for unusual 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",
        })
 
    # Check for high ratio of non-content tokens (HTML, markdown formatting)
    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" or "suspect"
    n_neighbors: int = 10,
) -> dict:
    """
    Analyze the distribution of embeddings to detect anomalous positioning.
 
    Poisoned documents may have embeddings that are:
    - Clustered around high-traffic query regions
    - Positioned as outliers relative to their supposed topic
    - Unusually close to embeddings from different topic categories
    """
    from sklearn.neighbors import NearestNeighbors
 
    # Compute nearest neighbor distances for each embedding
    nn = NearestNeighbors(n_neighbors=n_neighbors, metric='cosine')
    nn.fit(embeddings)
    distances, indices = nn.kneighbors(embeddings)
 
    # Separate baseline and suspect 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
 
    # For each suspect embedding, analyze its neighborhood
    suspect_indices = np.where(suspect_mask)[0]
 
    for idx in suspect_indices:
        neighbor_indices = indices[idx][1:]  # Exclude self
        neighbor_distances = distances[idx][1:]
 
        # What fraction of neighbors are baseline vs suspect?
        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))
 
        # Compute how this compares to baseline mean neighbor distance
        baseline_mean_dists = []
        for b_idx in np.where(baseline_mask)[0][:100]:  # Sample for efficiency
            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:
    """
    Analyze retrieval logs to detect RAG poisoning indicators.
 
    Poisoning manifests in retrieval patterns as:
    - New documents appearing with unusually high retrieval frequency
    - Documents retrieved for queries outside their expected topic scope
    - Sudden changes in which documents are retrieved for stable query types
    """
    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)
 
    # Identify suspiciously active documents
    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:
    """
    Assess the impact of identified poisoned documents.
 
    Determines how many queries were affected and what
    outputs were influenced by the poisoned content.
    """
    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],
    }