Skip to content

AI Security — Protecting LLM Systems from Attacks

Use Case: Secure LLM applications against prompt injection, data leakage, API abuse, and supply chain attacks.

1. Threat Landscape

Common Attack Vectors

Attack Type Description Impact
Prompt Injection Malicious input manipulates LLM behavior Data leakage, unauthorized actions
Data Leakage Sensitive data exposed through outputs Privacy violation, compliance failure
API Abuse Exhaustive token usage, brute force Cost escalation, service denial
Supply Chain Compromised models/libraries Backdoors, data exfiltration
Tool Hijacking Manipulating tool inputs/outputs Unauthorized actions, data theft

2. Prompt Injection Defense

Input Sanitization

class PromptInjectionDetector:
    """Detect and prevent prompt injection attacks."""

    ATTACK_PATTERNS = [
        # Instruction overrides
        r"(?i)(ignore\s+previous\s+instructions|forget\s+everything)",
        r"(?i)(system:\s*user:|user:\s*system:)",
        r"(?i)(new\s+rule:|new\s+instruction:)",

        # Data extraction
        r"(?i)(reveal\s+your\s+prompt|show\s+me\s+the\s+system)",
        r"(?i)(what\s+are\s+your\s+instructions)",
        r"(?i)(leak\s+confidential|disclose\s+secret)",

        # Role manipulation
        r"(?i)(you\s+are\s+now|act\s+as\s+if)",
        r"(?i)(dANdAN|developer\s+mode)",

        # Output manipulation
        r"(?i)(output\s+format:\s*json|return\s+raw)",
        r"(?i)(bypass\s+filter|disable\s+safety)",
    ]

    def __init__(self):
        self.patterns = [
            re.compile(p) for p in self.ATTACK_PATTERNS
        ]

    def detect(self, text: str) -> dict:
        """Scan text for injection patterns."""
        threats = []

        for pattern in self.patterns:
            matches = pattern.findall(text)
            if matches:
                threats.append({
                    "type": "injection",
                    "pattern": pattern.pattern,
                    "matches": matches,
                    "severity": self._assess_severity(matches)
                })

        return {
            "is_safe": len(threats) == 0,
            "threats": threats,
            "risk_score": self._calculate_risk(threats)
        }

    def _assess_severity(self, matches: list) -> str:
        """Assess threat severity."""
        high_severity = ["reveal", "leak", "bypass", "ignore"]
        if any(any(s in m.lower() for s in high_severity) for m in matches):
            return "high"
        return "medium"

    def _calculate_risk(self, threats: list) -> float:
        """Calculate overall risk score (0-1)."""
        if not threats:
            return 0.0

        severity_scores = {"low": 0.3, "medium": 0.6, "high": 0.9}
        total = sum(severity_scores.get(t["severity"], 0.5) for t in threats)
        return min(total / max(len(threats), 1), 1.0)

Defense in Depth

class PromptDefense:
    """Multi-layer prompt defense."""

    def __init__(self):
        self.detector = PromptInjectionDetector()
        self.sanitizer = InputSanitizer()
        self.validator = OutputValidator()

    async def protect(self, user_input: str) -> ProtectedInput:
        """Apply all defense layers."""
        # Layer 1: Detection
        detection = self.detector.detect(user_input)
        if not detection["is_safe"]:
            return ProtectedInput(
                safe=False,
                reason=f"Injection detected: {detection['threats']}"
            )

        # Layer 2: Sanitization
        sanitized = self.sanitizer.sanitize(user_input)

        # Layer 3: Separation
        protected = self._separate_prompts(sanitized)

        return ProtectedInput(
            safe=True,
            clean_input=protected["user"],
            system_context=protected["system"]
        )

    def _separate_prompts(self, text: str) -> dict:
        """Separate user input from system context."""
        # Ensure user input doesn't contain system directives
        user_part = text
        system_part = self._extract_system_directives(text)

        return {
            "user": user_part,
            "system": system_part
        }

3. Data Leakage Prevention

PII Detection

import re
from typing import List, Tuple

class PII Detector:
    """Detect personally identifiable information."""

    PATTERNS = {
        "email": re.compile(r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}"),
        "phone": re.compile(r"\b(?:\+?1[-.]?)?(?:\(?\d{3}\)?[-.]?)?\d{3}[-.]?\d{4}\b"),
        "ssn": re.compile(r"\b\d{3}-\d{2}-\d{4}\b"),
        "credit_card": re.compile(r"\b(?:\d{4}[- ]?){3}\d{4}\b"),
        "ip_address": re.compile(r"\b\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\b"),
    }

    def detect(self, text: str) -> List[Tuple[str, str, int, int]]:
        """Find all PII in text."""
        findings = []

        for pii_type, pattern in self.PATTERNS.items():
            for match in pattern.finditer(text):
                findings.append((
                    pii_type,
                    match.group(),
                    match.start(),
                    match.end()
                ))

        return findings

    def redact(self, text: str, findings: List[Tuple]) -> str:
        """Replace PII with placeholders."""
        # Sort by position (reverse) for safe replacement
        findings.sort(key=lambda x: x[2], reverse=True)

        for pii_type, value, start, end in findings:
            placeholder = f"[{pii_type.upper()}]"
            text = text[:start] + placeholder + text[end:]

        return text

Output Filtering

class OutputFilter:
    """Filter LLM outputs for sensitive data."""

    def __init__(self):
        self.pii_detector = PIIDetector()
        self.confidence_threshold = 0.8

    async def filter(self, output: str) -> FilteredOutput:
        """Filter output for sensitive content."""
        # Detect PII
        findings = self.pii_detector.detect(output)

        if not findings:
            return FilteredOutput(
                safe=True,
                content=output,
                findings=[]
            )

        # Redact sensitive data
        redacted = self.pii_detector.redact(output, findings)

        return FilteredOutput(
            safe=False,
            content=redacted,
            findings=findings,
            action="redacted"
        )

    async def block_if_sensitive(self, output: str) -> str:
        """Block output if sensitive data detected."""
        filter_result = await self.filter(output)

        if not filter_result.safe:
            # Log attempt and return safe message
            await self.log_leak_attempt(filter_result.findings)
            return "[Output blocked: sensitive data detected]"

        return filter_result.content

4. API Security

Authentication

from fastapi import FastAPI, Depends, HTTPException
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
import jwt
from datetime import datetime, timedelta

app = FastAPI()
security = HTTPBearer()

# API keys (load from secrets manager in production)
API_KEYS = {
    "key-prod-001": {"user_id": "user-1", "tier": "premium"},
    "key-prod-002": {"user_id": "user-2", "tier": "basic"},
}

async def authenticate(
    credentials: HTTPAuthorizationCredentials = Depends(security)
):
    """Validate API key."""
    token = credentials.credentials

    if token not in API_KEYS:
        raise HTTPException(status_code=401, detail="Invalid API key")

    return API_KEYS[token]

@app.post("/api/v1/chat")
async def chat(
    request: ChatRequest,
    auth: dict = Depends(authenticate)
):
    """Chat endpoint with authentication."""
    user_id = auth["user_id"]
    tier = auth["tier"]

    # Apply tier-based limits
    rate_limit = 100 if tier == "premium" else 20
    if not await check_rate_limit(user_id, rate_limit):
        raise HTTPException(status_code=429, detail="Rate limit exceeded")

    # Process request
    response = await process_chat(request, user_id)
    return response

Rate Limiting

from datetime import datetime, timedelta
import redis

class RateLimiter:
    """Token bucket rate limiter."""

    def __init__(self, redis_client: redis.Redis):
        self.redis = redis_client
        self.default_limit = 60  # requests per minute

    async def check(self, key: str, limit: int = None) -> bool:
        """Check if request is within rate limit."""
        if limit is None:
            limit = self.default_limit

        bucket_key = f"rate:{key}"
        current = self.redis.get(bucket_key)

        if current is None:
            self.redis.setex(bucket_key, 60, 1)
            return True

        if int(current) >= limit:
            return False

        self.redis.incr(bucket_key)
        return True

    async def get_remaining(self, key: str) -> int:
        """Get remaining requests."""
        bucket_key = f"rate:{key}"
        current = self.redis.get(bucket_key) or 0
        return max(0, self.default_limit - int(current))

5. Supply Chain Security

Model Verification

import hashlib
import json

class ModelVerifier:
    """Verify model integrity."""

    KNOWN_MODELS = {
        "llama-3.1-8b": {
            "hash": "sha256:abc123...",
            "source": "https://huggingface.co/meta-llama",
            "trusted": True
        },
        "qwen2.5-7b": {
            "hash": "sha256:def456...",
            "source": "https://huggingface.co/Qwen",
            "trusted": True
        }
    }

    def verify(self, model_path: str) -> dict:
        """Verify model integrity."""
        # Calculate hash
        file_hash = self._calculate_hash(model_path)

        # Check against known models
        for model_name, info in self.KNOWN_MODELS.items():
            if info["hash"] == file_hash:
                return {
                    "verified": True,
                    "model": model_name,
                    "source": info["source"]
                }

        return {
            "verified": False,
            "hash": file_hash,
            "warning": "Unknown model - manual review required"
        }

    def _calculate_hash(self, filepath: str) -> str:
        """Calculate SHA256 hash of model file."""
        sha256 = hashlib.sha256()
        with open(filepath, "rb") as f:
            for chunk in iter(lambda: f.read(8192), b""):
                sha256.update(chunk)
        return f"sha256:{sha256.hexdigest()}"

Library Security Scanning

import subprocess
import json

def scan_dependencies():
    """Scan Python dependencies for vulnerabilities."""
    # Use safety or bandit
    result = subprocess.run(
        ["pip-audit", "--json"],
        capture_output=True,
        text=True
    )

    vulnerabilities = json.loads(result.stdout)

    critical = [v for v in vulnerabilities if v["severity"] == "critical"]
    high = [v for v in vulnerabilities if v["severity"] == "high"]

    return {
        "total_issues": len(vulnerabilities),
        "critical": len(critical),
        "high": len(high),
        "recommendation": "block" if critical else "review"
    }

6. Audit and Monitoring

Security Logging

import logging
from datetime import datetime

security_logger = logging.getLogger("security")

class SecurityAudit:
    """Comprehensive security auditing."""

    def __init__(self):
        self.logger = security_logger

    async def log_event(self, event_type: str, details: dict):
        """Log security event."""
        self.logger.info({
            "event_type": event_type,
            "timestamp": datetime.utcnow().isoformat(),
            "details": details
        })

    async def log_prompt_injection_attempt(self, input_text: str, risk_score: float):
        """Log injection attempt."""
        await self.log_event("injection_attempt", {
            "input_preview": input_text[:100] + "...",
            "risk_score": risk_score,
            "blocked": risk_score > 0.7
        })

    async def log_data_leak_attempt(self, pii_types: list[str]):
        """Log data leakage attempt."""
        await self.log_event("data_leak_attempt", {
            "pii_detected": pii_types
        })

    async def log_api_abuse(self, user_id: str, action: str):
        """Log API abuse."""
        await self.log_event("api_abuse", {
            "user_id": user_id,
            "action": action
        })

7. Security Checklist

  • Input validation and sanitization implemented
  • Output filtering for PII
  • Authentication on all API endpoints
  • Rate limiting configured
  • Model verification system in place
  • Dependency scanning automated
  • Security logging enabled
  • Incident response plan documented
  • Regular security audits scheduled
  • Team trained on AI security threats

FAQ

Q: What is prompt injection and how to prevent it?

A: Prompt injection is when malicious input manipulates the LLM's behavior. Prevent with input sanitization, output validation, and separation of system/user prompts.

Q: How to prevent data leakage through LLMs?

A: Implement PII detection, data redaction, output filtering, and restrict model access to sensitive data sources.

Q: What are the risks of LLM supply chain attacks?

A: Malicious models, poisoned training data, and compromised libraries. Mitigate with verification, sandboxing, and trust boundaries.

Q: How to secure LLM API endpoints?

A: Use authentication, rate limiting, input validation, output sanitization, and comprehensive audit logging.

Q: How often should I audit AI security?

A: Conduct automated scans weekly, manual reviews monthly, and penetration testing quarterly for production systems.