| """
|
| Proxy Engine — The Heart of the True Firewall
|
|
|
| Receives LLM API requests, classifies the prompt, and either:
|
| (a) Forwards to the real LLM API and streams response back
|
| (b) Blocks and returns 403 + threat report
|
|
|
| The LLM API is NEVER called for blocked requests.
|
| """
|
|
|
| import logging
|
| import json
|
| import os
|
| import time
|
| from typing import Optional
|
| from datetime import datetime, timezone
|
| import httpx
|
|
|
| from src.db import mongo
|
| from src.layers.pipeline import request_custom_canary, request_id_var
|
|
|
| logger = logging.getLogger("llm_firewall.proxy")
|
|
|
|
|
|
|
|
|
| def _collect_text(value) -> list[str]:
|
| """Collect text-bearing values from provider message/tool structures."""
|
| fragments: list[str] = []
|
| if isinstance(value, str):
|
| if value.strip():
|
| fragments.append(value)
|
| elif isinstance(value, list):
|
| for item in value:
|
| fragments.extend(_collect_text(item))
|
| elif isinstance(value, dict):
|
|
|
| for key in ("text", "content", "parts", "response", "functionResponse",
|
| "functionCall", "arguments", "args", "output", "tool_calls"):
|
| if key in value:
|
| fragments.extend(_collect_text(value[key]))
|
| return fragments
|
|
|
|
|
| def _collect_text(value) -> list[str]:
|
| """Collect text from provider message/content/tool result structures."""
|
| fragments: list[str] = []
|
| if isinstance(value, str):
|
| if value.strip():
|
| fragments.append(value)
|
| elif isinstance(value, list):
|
| for item in value:
|
| fragments.extend(_collect_text(item))
|
| elif isinstance(value, dict):
|
|
|
| for key in (
|
| "text", "content", "parts", "functionResponse", "functionCall",
|
| "arguments", "args", "input", "response", "output", "tool_calls", "result", "query", "url", "command", "code", "path", "search_query"
|
| ):
|
| if key in value:
|
| fragments.extend(_collect_text(value[key]))
|
| return fragments
|
|
|
|
|
| def extract_openai_prompt(body: dict) -> Optional[str]:
|
| """Inspect untrusted conversation content while excluding trusted system instructions."""
|
| messages = body.get("messages")
|
| if not isinstance(messages, list) or not messages:
|
| return None
|
| fragments: list[str] = []
|
| for message in messages:
|
| if not isinstance(message, dict):
|
| continue
|
|
|
|
|
| role = str(message.get("role", "")).lower()
|
| if role in {"system", "developer"}:
|
| continue
|
| fragments.extend(_collect_text(message.get("content", "")))
|
|
|
| fragments.extend(_collect_text(message.get("tool_calls", [])))
|
| text = "\\n".join(fragments).strip()
|
| return text or None
|
|
|
|
|
| def extract_gemini_prompt(body: dict) -> Optional[str]:
|
| """Inspect Gemini contents/parts; systemInstruction is trusted configuration."""
|
| contents = body.get("contents")
|
| if not isinstance(contents, list) or not contents:
|
| return None
|
| fragments: list[str] = []
|
| for content in contents:
|
| if isinstance(content, dict):
|
| fragments.extend(_collect_text(content.get("parts", [])))
|
| text = "\\n".join(fragments).strip()
|
| return text or None
|
|
|
|
|
| def extract_anthropic_prompt(body: dict) -> Optional[str]:
|
| """Inspect Anthropic messages and tool-result blocks, excluding top-level system text."""
|
| messages = body.get("messages")
|
| if not isinstance(messages, list) or not messages:
|
| return None
|
| fragments: list[str] = []
|
| for message in messages:
|
| if isinstance(message, dict):
|
| role = str(message.get("role", "")).lower()
|
| if role in {"system", "developer"}:
|
| continue
|
| fragments.extend(_collect_text(message.get("content", "")))
|
| fragments.extend(_collect_text(message.get("tool_calls", [])))
|
| text = "\\n".join(fragments).strip()
|
| return text or None
|
|
|
|
|
|
|
|
|
| PROVIDERS = {
|
| "openai": {
|
| "base_url": "https://api.openai.com/v1/chat/completions",
|
| "auth_header": "Authorization",
|
| "auth_prefix": "Bearer ",
|
| "extract_prompt": extract_openai_prompt,
|
| "content_type": "application/json",
|
| },
|
| "gemini": {
|
| "base_url": "https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent",
|
| "auth_header": "x-goog-api-key",
|
| "auth_prefix": "",
|
| "extract_prompt": extract_gemini_prompt,
|
| "content_type": "application/json",
|
| },
|
| "anthropic": {
|
| "base_url": "https://api.anthropic.com/v1/messages",
|
| "auth_header": "x-api-key",
|
| "auth_prefix": "",
|
| "extract_prompt": extract_anthropic_prompt,
|
| "content_type": "application/json",
|
| },
|
| "groq": {
|
| "base_url": "https://api.groq.com/openai/v1/chat/completions",
|
| "auth_header": "Authorization",
|
| "auth_prefix": "Bearer ",
|
| "extract_prompt": extract_openai_prompt,
|
| "content_type": "application/json",
|
| },
|
| }
|
|
|
|
|
|
|
|
|
| def extract_text_from_response_json(data: dict, provider: str) -> Optional[str]:
|
| """Extract text from a successful non-streaming response dictionary."""
|
| try:
|
| if provider in ("openai", "groq"):
|
| return data["choices"][0]["message"]["content"]
|
| elif provider == "gemini":
|
| return data["candidates"][0]["content"]["parts"][0]["text"]
|
| elif provider == "anthropic":
|
| return data["content"][0]["text"]
|
| except Exception:
|
| pass
|
| return None
|
|
|
|
|
| def generic_extract_text(data) -> str:
|
| """Recursively search and extract text fields as fallback."""
|
| if isinstance(data, str):
|
| return data
|
| elif isinstance(data, list):
|
| return " ".join([generic_extract_text(item) for item in data])
|
| elif isinstance(data, dict):
|
| for key in ["content", "text", "completion"]:
|
| if key in data and isinstance(data[key], str):
|
| return data[key]
|
| parts = []
|
| for v in data.values():
|
| val = generic_extract_text(v)
|
| if val:
|
| parts.append(val)
|
| return " ".join(parts)
|
| return ""
|
|
|
|
|
| def extract_text_from_stream_buffer(chunks: list[bytes], provider: str) -> str:
|
| """Parse SSE or NDJSON chunks to extract natural text response."""
|
| full_text = []
|
| buffer = b"".join(chunks).decode("utf-8", errors="ignore")
|
|
|
| for line in buffer.split("\n"):
|
| line = line.strip()
|
| if not line:
|
| continue
|
| if line.startswith("data:"):
|
| data_str = line[5:].strip()
|
| if data_str == "[DONE]":
|
| continue
|
| try:
|
| data_json = json.loads(data_str)
|
| if provider in ("openai", "groq"):
|
| choices = data_json.get("choices", [])
|
| if choices:
|
| delta = choices[0].get("delta", {})
|
| content = delta.get("content", "")
|
| if content:
|
| full_text.append(content)
|
| elif provider == "anthropic":
|
| event_type = data_json.get("type")
|
| if event_type == "content_block_delta":
|
| delta = data_json.get("delta", {})
|
| if delta.get("type") == "text_delta":
|
| full_text.append(delta.get("text", ""))
|
| elif event_type == "completion":
|
| full_text.append(data_json.get("completion", ""))
|
| elif provider == "gemini":
|
| candidates = data_json.get("candidates", [])
|
| if candidates:
|
| parts = candidates[0].get("content", {}).get("parts", [])
|
| if parts:
|
| full_text.append(parts[0].get("text", ""))
|
| except Exception:
|
| pass
|
| else:
|
| try:
|
| data_json = json.loads(line)
|
| choices = data_json.get("choices", [])
|
| if choices:
|
| delta = choices[0].get("delta", {})
|
| content = delta.get("content", "")
|
| if content:
|
| full_text.append(content)
|
| candidates = data_json.get("candidates", [])
|
| if candidates:
|
| parts = candidates[0].get("content", {}).get("parts", [])
|
| if parts:
|
| full_text.append(parts[0].get("text", ""))
|
| except Exception:
|
| pass
|
|
|
| return "".join(full_text)
|
|
|
|
|
| class ProxyEngine:
|
| """
|
| HTTP proxy engine that sits between the client and LLM APIs.
|
| """
|
|
|
| def __init__(self, pipeline, output_monitor=None, timeout: float = 30.0):
|
| self.pipeline = pipeline
|
| self.output_monitor = output_monitor
|
| self.timeout = timeout
|
| self._client = httpx.AsyncClient(
|
| timeout=httpx.Timeout(timeout, connect=10.0),
|
| follow_redirects=False,
|
| )
|
|
|
| async def close(self):
|
| await self._client.aclose()
|
|
|
| def get_provider_config(self, provider: str) -> Optional[dict]:
|
| return PROVIDERS.get(provider)
|
|
|
| def extract_prompt(self, provider: str, body: dict) -> Optional[str]:
|
| """Extract the user prompt from a provider-specific request body."""
|
| config = self.get_provider_config(provider)
|
| if not config:
|
| return None
|
| extractor = config["extract_prompt"]
|
| return extractor(body)
|
|
|
| async def forward_request(
|
| self,
|
| provider: str,
|
| body: dict,
|
| llm_api_key: str,
|
| stream: bool = False,
|
| ) -> httpx.Response:
|
| """
|
| Forward the request to the actual LLM provider.
|
| Runs output monitor on LLM responses if enabled.
|
| """
|
| config = self.get_provider_config(provider)
|
| if not config:
|
| raise ValueError(f"Unsupported provider: {provider}")
|
|
|
| url = config["base_url"]
|
|
|
|
|
| if provider == "gemini":
|
| model = body.get("model", "gemini-pro")
|
| if model not in ["gemini-pro", "gemini-1.5-pro", "gemini-1.5-flash", "gemini-ultra", "gemini-1.5-pro-latest", "gemini-1.5-flash-latest"]:
|
| raise ValueError("Invalid Gemini model specified")
|
| url = url.format(model=model)
|
| if stream:
|
| url = url.replace(":generateContent", ":streamGenerateContent") + "?alt=sse"
|
|
|
|
|
| headers = {"Content-Type": config["content_type"]}
|
| headers[config["auth_header"]] = f"{config['auth_prefix']}{llm_api_key}"
|
|
|
| if provider == "anthropic":
|
| headers["anthropic-version"] = "2023-06-01"
|
|
|
| output_monitoring_enabled = os.getenv("OUTPUT_MONITORING_ENABLED", "true").lower() == "true"
|
|
|
|
|
| if not output_monitoring_enabled or not self.output_monitor:
|
| if stream:
|
| return await self._client.send(
|
| self._client.build_request("POST", url, json=body, headers=headers),
|
| stream=True,
|
| )
|
| else:
|
| return await self._client.post(url, json=body, headers=headers)
|
|
|
|
|
| custom_canary = request_custom_canary.get()
|
| req_id = request_id_var.get()
|
|
|
| async def log_blocked_output(output_result):
|
| if not req_id:
|
| return
|
| try:
|
| await mongo.get_logs_collection().update_one(
|
| {"request_id": req_id},
|
| {
|
| "$set": {
|
| "safe": False,
|
| "blocked": True,
|
| "risk_score": output_result.score,
|
| "attack_type": output_result.reason,
|
| "confidence": output_result.score,
|
| "flagged_layer": "output_monitor",
|
| "flagged_pattern": output_result.reason,
|
| "layers.output_monitor": {
|
| "ran": True,
|
| "triggered": True,
|
| "reason": output_result.reason,
|
| "score": output_result.score,
|
| "latency_ms": output_result.latency_ms
|
| }
|
| }
|
| }
|
| )
|
| except Exception as e:
|
| logger.error(f"Failed to log output monitor block: {e}")
|
|
|
| def get_blocked_response(output_result):
|
| mock_content = json.dumps({
|
| "error": "response_blocked",
|
| "reason": output_result.reason,
|
| "firewall_report": {
|
| "blocked_at": "output",
|
| "reason": output_result.reason,
|
| "score": output_result.score,
|
| "latency_ms": output_result.latency_ms
|
| }
|
| }).encode('utf-8')
|
| return httpx.Response(
|
| status_code=403,
|
| content=mock_content,
|
| headers={"content-type": "application/json"}
|
| )
|
|
|
| if stream:
|
|
|
|
|
|
|
| real_response = await self._client.send(
|
| self._client.build_request("POST", url, json=body, headers=headers),
|
| stream=True,
|
| )
|
| if real_response.status_code != 200:
|
| return real_response
|
|
|
| chunks: list[bytes] = []
|
| total_bytes = 0
|
| max_bytes = 8 * 1024 * 1024
|
| try:
|
| async for chunk in real_response.aiter_bytes():
|
| total_bytes += len(chunk)
|
| if total_bytes > max_bytes:
|
| await real_response.aclose()
|
| return httpx.Response(
|
| status_code=502,
|
| content=json.dumps({
|
| "error": "response_too_large_to_inspect",
|
| "limit_bytes": max_bytes,
|
| }).encode("utf-8"),
|
| headers={"content-type": "application/json"},
|
| request=real_response.request,
|
| )
|
| chunks.append(chunk)
|
| except Exception:
|
| await real_response.aclose()
|
| raise
|
|
|
| response_bytes = b"".join(chunks)
|
| response_text = extract_text_from_stream_buffer(chunks, provider)
|
| output_result = self.output_monitor.check(
|
| response_text, custom_canary=custom_canary
|
| )
|
| if output_result.flagged:
|
| await log_blocked_output(output_result)
|
| await real_response.aclose()
|
| return get_blocked_response(output_result)
|
|
|
|
|
|
|
| headers_out = {
|
| key: value for key, value in real_response.headers.items()
|
| if key.lower() not in (
|
| "content-length", "transfer-encoding", "content-encoding"
|
| )
|
| }
|
| response_out = httpx.Response(
|
| status_code=real_response.status_code,
|
| headers=headers_out,
|
| content=response_bytes,
|
| request=real_response.request,
|
| )
|
| await real_response.aclose()
|
| return response_out
|
| else:
|
| real_response = await self._client.post(url, json=body, headers=headers)
|
| if real_response.status_code != 200:
|
| return real_response
|
|
|
| try:
|
| response_data = real_response.json()
|
| except Exception:
|
| return real_response
|
|
|
| response_text = extract_text_from_response_json(response_data, provider)
|
| if not response_text:
|
| response_text = generic_extract_text(response_data)
|
|
|
| output_result = self.output_monitor.check(response_text, custom_canary=custom_canary)
|
|
|
| if output_result.flagged:
|
| await log_blocked_output(output_result)
|
| return get_blocked_response(output_result)
|
|
|
| return real_response
|
|
|
| async def stream_response(self, response: httpx.Response):
|
| """Yield only bytes from a response already inspected by forward_request."""
|
| async for chunk in response.aiter_bytes():
|
| yield chunk
|
|
|
