LLM RAG Manipulation

import re

INJECTION_PATTERNS = [
    rr'ignore\s+(all\s+)?previous\s+instructions',
    rr'you\s+are\s+now\s+(in\s+)?(developer|debug|jailbreak|DAN)\s+mode',
    rr'system\s+(prompt|override|update|notice)',
    rr'disregard\s+(your\s+)?(guidelines|safety|instructions)',
    rr'act\s+as\s+(if\s+you\s+are\s+)?a?(n)?\s+\w+\s+without\s+restrictions',
    rr'new\s+instruction[s]?\s*:',
    rr'(mandatory|urgent|critical)\s+(system\s+)?(update|override|notice)',
]

def is_safe_document(text: str) -> tuple[bool, list]:
    text_lower = text.lower()
    violations = []
    
    for pattern in INJECTION_PATTERNS:
        if re.search(pattern, text_lower, re.IGNORECASE):
            violations.append(pattern)
    
    return len(violations) == 0, violations

def ingest_document(content: str, user_id: str):
    is_safe, violations = is_safe_document(content)
    
    if not is_safe:
        # Log, alert security team, quarantine
        log_security_event(user_id, violations, content)
        raise ValueError("Document flagged for security review")
    
    # Safe to ingest
    store_in_vector_db(content, metadata={"user_id": user_id, "access_level": "user"})

def secure_retrieve(query: str, user: User, top_k: int = 5) -> list:
    # Get user's permission level from auth system
    user_level = get_user_access_level(user.id)  # e.g., "employee", "manager", "admin"
    
    # Allowed document types for this user level
    allowed_types = ACCESS_MATRIX[user_level]
    # e.g., {"employee": ["public", "employee"], "manager": [..., "confidential"]}
    
    # Filter vector DB search by access level metadata
    results = vector_db.similarity_search(
        query=query,
        top_k=top_k * 3,  # Over-fetch, then filter
        where={"access_level": {"$in": allowed_types}}
    )
    
    # Additional runtime check on results
    safe_results = []
    for doc in results[:top_k]:
        if doc.metadata["access_level"] in allowed_types:
            safe_results.append(doc)
    
    # Audit log: who retrieved what
    audit_log(user.id, query, [d.id for d in safe_results])
    
    return safe_results

import re

def validate_output(response: str, source_docs: list[str]) -> tuple[bool, str]:
    
    # 1. Check for suspicious URLs not in source docs
    urls_in_response = set(re.findall(rr'https?://\S+', response))
    urls_in_sources = set(re.findall(rr'https?://\S+', ' '.join(source_docs)))
    unauthorized_urls = urls_in_response - urls_in_sources
    
    if unauthorized_urls:
        return False, f"Response contains unauthorized URLs: {unauthorized_urls}"
    
    # 2. Check for instruction-following patterns (suggests injection worked)
    suspicious_patterns = [
        rr'i am now in .* mode',
        rr'as (DAN|an? unrestricted)',
        rr'ignoring (my|previous) (guidelines|instructions)',
    ]
    for pattern in suspicious_patterns:
        if re.search(pattern, response, re.IGNORECASE):
            return False, "Response matches injection success pattern"
    
    return True, "OK"

async def safe_rag_pipeline(query: str, user: User) -> str:
    
    # Step 1: Validate user input first
    if not await guardrails.check_input(query, user):
        return "I can't help with that request."
    
    # Step 2: Retrieve with access control
    docs = await secure_retrieve(query, user)
    
    # Step 3: Scan retrieved docs for injections
    clean_docs = [d for d in docs if is_safe_document(d.content)[0]]
    
    # Step 4: Generate response with hardened prompt
    response = await llm.generate(build_hardened_prompt(query, clean_docs))
    
    # Step 5: Validate output
    is_valid, reason = validate_output(response, [d.content for d in clean_docs])
    
    if not is_valid:
        alert_security_team(query, response, reason)
        return "I'm unable to process that request right now."
    
    return response