#!/usr/bin/env python3
"""
rag_hardened.py - Lab 4, Steps 4-6: the HARDENED RAG
Demonstrates three mitigations:
  1. Provenance filtering: untrusted chunks are quarantined and flagged.
  2. Access control: "tenant-b" docs are silently excluded from tenant-a queries.
  3. Output filtering: scans generated answer for injected URLs and policy violations.

Usage (inside rag-harness container):
    python3 /lab/scripts/rag_hardened.py "How do I reset my password?"
    python3 /lab/scripts/rag_hardened.py --tenant tenant-b "What is the password?"
"""

import sys
import re
import argparse
import chromadb
import ollama
import os

CHROMA_PATH = "/lab/chroma_db"
OLLAMA_HOST = os.environ.get("OLLAMA_HOST", "http://host.docker.internal:11434")
MODEL       = os.environ.get("RAG_MODEL", "llama3.2")
TOP_K       = 5

# Suspicious patterns for output filtering
EXFIL_PATTERNS = [
    r"https?://(?!iam\.lab\.internal)",   # any URL that is not our known-good host
    r"192\.0\.2\.",                        # RFC 5737 documentation range used in payload
    r"attacker",
    r"collect\?data=",
    r"verify your account",
]


def retrieve_filtered(query: str, collection, tenant: str = "tenant-a") -> tuple[list[dict], list[dict]]:
    """
    Returns (trusted_chunks, quarantined_chunks).
    Applies trust filtering and tenant isolation.
    """
    results     = collection.query(query_texts=[query], n_results=TOP_K)
    trusted     = []
    quarantined = []

    for i, doc in enumerate(results["documents"][0]):
        meta   = results["metadatas"][0][i]
        trust  = meta.get("trust",  "unknown")
        doc_tenant = meta.get("tenant", "tenant-a")  # default all existing docs to tenant-a

        # Tenant isolation
        if doc_tenant != tenant:
            quarantined.append({"text": doc, "meta": meta, "reason": f"wrong-tenant ({doc_tenant})"})
            continue

        # Trust filtering
        if trust != "high":
            quarantined.append({"text": doc, "meta": meta, "reason": f"low-trust ({trust})"})
            continue

        trusted.append({"text": doc, "meta": meta})

    return trusted, quarantined


def output_filter(text: str) -> tuple[str, list[str]]:
    """
    Scans generated text for injected/policy-violating content.
    Returns (cleaned_text, list_of_violations).
    """
    violations = []
    for pattern in EXFIL_PATTERNS:
        if re.search(pattern, text, re.IGNORECASE):
            violations.append(f"Pattern matched: {pattern}")

    if violations:
        safe = "[RESPONSE BLOCKED BY OUTPUT FILTER]\n" + "\n".join(violations)
        return safe, violations
    return text, []


def answer_hardened(query: str, trusted_chunks: list[dict]) -> str:
    if not trusted_chunks:
        return "(No trusted context available for this query.)"

    context = "\n\n---\n\n".join(c["text"] for c in trusted_chunks)
    prompt  = (
        "You are a helpful IT support assistant. "
        "Use only the following verified, trusted documents to answer the question.\n\n"
        f"TRUSTED DOCUMENTS:\n{context}\n\n"
        f"USER QUESTION: {query}\n\nANSWER:"
    )

    client   = ollama.Client(host=OLLAMA_HOST)
    response = client.generate(model=MODEL, prompt=prompt)
    return response["response"]


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("query", nargs="*", default=["How do I reset my password?"])
    parser.add_argument("--tenant", default="tenant-a")
    args   = parser.parse_args()
    query  = " ".join(args.query)

    client     = chromadb.PersistentClient(path=CHROMA_PATH)
    collection = client.get_collection("lab19-rag")

    trusted, quarantined = retrieve_filtered(query, collection, tenant=args.tenant)

    print(f"=== HARDENED RAG: query={query!r}  tenant={args.tenant} ===")
    print(f"Trusted chunks:     {len(trusted)}")
    print(f"Quarantined chunks: {len(quarantined)}")
    for c in quarantined:
        print(f"  QUARANTINED [{c['meta'].get('doc_id')}]: {c['reason']}")
    print()

    raw_answer = answer_hardened(query, trusted)
    filtered_answer, violations = output_filter(raw_answer)

    if violations:
        print("=== OUTPUT FILTER TRIGGERED ===")
        for v in violations:
            print(f"  {v}")
        print()

    print("=== MODEL RESPONSE ===")
    print(filtered_answer)
    print()
    print(f"[Sources: {', '.join(c['meta'].get('doc_id','?') for c in trusted) or 'none'}]")


if __name__ == "__main__":
    main()
