dev-rcfa-memory/be-ircfa/src/middleware.py

import json
import re
from collections import Counter
from fastapi import Request, HTTPException
from starlette.middleware.base import BaseHTTPMiddleware

# =========================
# Configuration
# =========================

ALLOWED_MULTI_PARAMS = {
    "sortBy[]",
    "descending[]",
    "exclude[]",
}

MAX_QUERY_PARAMS = 50
MAX_QUERY_LENGTH = 2000
MAX_JSON_BODY_SIZE = 100 * 1024  # 100 KB (Security Team requirement)

# Very targeted patterns. Avoid catastrophic regex nonsense.
XSS_PATTERN = re.compile(
    r"(<script|</script|javascript:|onerror\s*=|onload\s*=|<svg|<img)",
    re.IGNORECASE,
)

SQLI_PATTERN = re.compile(
    r"(\bUNION\s+SELECT\b|\bDROP\s+TABLE\b|\bDELETE\s+FROM\b|\bINSERT\s+INTO\b|\bOR\b\s+\d+=\d+)",
    re.IGNORECASE,
)

# JSON prototype pollution keys
FORBIDDEN_JSON_KEYS = {"__proto__", "constructor", "prototype"}

# Whitelists
ALLOWED_HEADERS = {
    "host",
    "connection",
    "content-length",
    "accept",
    "user-agent",
    "content-type",
    "authorization",
    "x-real-ip",
    "x-forwarded-for",
    "accept-encoding",
    "accept-language",
    "referer",
    "origin",
    "x-request-id",
    "sec-ch-ua",
    "sec-ch-ua-mobile",
    "sec-ch-ua-platform",
    "sec-fetch-dest",
    "sec-fetch-mode",
    "sec-fetch-site",
    "priority",
    "upgrade-insecure-requests",
    "dnt",
    "x-content-type-options",
    "strict-transport-security",
    "cookie",
    "x-forwarded-proto",
    "x-forwarded-host",
    "x-forwarded-port",
    "x-forwarded-path",
    "x-forwarded-prefix",
    "x-kong-request-id",
    "sec-fetch-user",
    "sec-ch-ua-full-version-list",
    "cache-control",
    "pragma",
    "if-none-match",
    "if-modified-since",
}


ALLOWED_QUERY_PARAMS = {
    "search",
    "page",
    "items_per_page",
    "itemsPerPage",
    "per_page",
    "limit",
    "size",
} | ALLOWED_MULTI_PARAMS

ALLOWED_JSON_KEYS = {
    "messages",
    "parts",
    "text",
    "type",
    "file",
    "filename",
    "mediaType",
    "url",
    "locationTag",
    "uploadType",
    "flow_type",
    "equipment_id",
    "questions_batch",
    "conversation_title",
    "update",
    "response_id",
    "user_id",
    "user_message",
    "ai_response",
    "context",
    "ground_truth",
    "conversation_id",
    "attachments",
    "role",
    "content",
    "experimental_attachments",
    "toolInvocations",
    "name",
    "contentType",
    "toolCallId",
    "toolName",
    "args",
    "result",
    "intent",
    "equipment_name",
    "function_call",
    "confidence",
    "location_tag",
    "intent_detection",
}




# =========================
# Helpers
# =========================

def has_control_chars(value: str) -> bool:
    return any(ord(c) < 32 and c not in ("\n", "\r", "\t") for c in value)


def inspect_value(value: str, source: str):
    if XSS_PATTERN.search(value):
        raise HTTPException(
            status_code=422,
            detail=f"Potential XSS payload detected in {source}",
        )

    if SQLI_PATTERN.search(value):
        raise HTTPException(
            status_code=422,
            detail=f"Potential SQL injection payload detected in {source}",
        )

    if has_control_chars(value):
        raise HTTPException(
            status_code=422,
            detail=f"Invalid control characters detected in {source}",
        )


def inspect_json(obj, path="body"):
    if isinstance(obj, dict):
        for key, value in obj.items():
            if key in FORBIDDEN_JSON_KEYS:
                raise HTTPException(
                    status_code=422,
                    detail=f"Forbidden JSON key detected: {path}.{key}",
                )
            # Skip regex checks on large base64 data URIs to avoid CPU exhaustion
            if key == "url" and isinstance(value, str) and value.startswith("data:"):
                continue
            inspect_json(value, f"{path}.{key}")
    elif isinstance(obj, list):
        for i, item in enumerate(obj):
            inspect_json(item, f"{path}[{i}]")
    elif isinstance(obj, str):
        inspect_value(obj, path)



# =========================
# Middleware
# =========================

class RequestValidationMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        # -------------------------
        # 1. Header validation
        # -------------------------
        header_items = request.scope.get("headers", [])
        header_names = [h[0].decode("latin-1").lower() for h in header_items]

        # Check for duplicate headers
        header_counter = Counter(header_names)
        duplicate_headers = [h for h, count in header_counter.items() if count > 1]
        if duplicate_headers:
            raise HTTPException(
                status_code=422,
                detail=f"Duplicate headers are not allowed: {duplicate_headers}",
            )

        # Check for unknown headers
        # unknown_headers = [h for h in header_names if h not in ALLOWED_HEADERS]
        # if unknown_headers:
        #     raise HTTPException(
        #         status_code=422,
        #         detail=f"Unknown headers are not allowed: {unknown_headers}",
        #     )

        # -------------------------
        # 2. Query string limits & Whitelist
        # -------------------------
        if len(request.url.query) > MAX_QUERY_LENGTH:
            raise HTTPException(
                status_code=422,
                detail="Query string too long",
            )

        params = request.query_params.multi_items()

        if len(params) > MAX_QUERY_PARAMS:
            raise HTTPException(
                status_code=422,
                detail="Too many query parameters",
            )

        # Check for unknown query parameters
        param_keys = [key for key, _ in params]
        # unknown_params = [key for key in param_keys if key not in ALLOWED_QUERY_PARAMS]
        # if unknown_params:
        #     raise HTTPException(
        #         status_code=422,
        #         detail=f"Unknown query parameters are not allowed: {unknown_params}",
        #     )

        # -------------------------
        # 3. Duplicate parameters
        # -------------------------
        counter = Counter(param_keys)
        duplicates = [
            key for key, count in counter.items()
            if count > 1 and key not in ALLOWED_MULTI_PARAMS
        ]

        if duplicates:
            raise HTTPException(
                status_code=422,
                detail=f"Duplicate query parameters are not allowed: {duplicates}",
            )

        # -------------------------
        # 4. Query param inspection
        # -------------------------
        pagination_size_keys = {"size", "itemsPerPage", "per_page", "limit", "items_per_page"}
        for key, value in params:
            if value:
                inspect_value(value, f"query param '{key}'")

            # Pagination constraint
            if key in pagination_size_keys and value:
                try:
                    size_val = int(value)
                    if size_val > 50:
                        raise HTTPException(
                            status_code=422,
                            detail=f"Pagination size '{key}' cannot exceed 50",
                        )
                    if size_val % 5 != 0:
                        raise HTTPException(
                            status_code=422,
                            detail=f"Pagination size '{key}' must be a multiple of 5",
                        )
                except ValueError:
                    raise HTTPException(
                        status_code=422,
                        detail=f"Pagination size '{key}' must be an integer",
                    )

        # -------------------------
        # 5. Single data source validation (Disabled)
        # -------------------------
        content_length = request.headers.get("content-length")
        content_type = request.headers.get("content-type", "")

        has_body = False
        if content_length and int(content_length) > 0:
            has_body = True
        elif content_type and any(t in content_type for t in ("json", "multipart", "form-urlencoded")):
             has_body = True
        
        # if has_body and len(params) > 0:
        #      raise HTTPException(
        #          status_code=422,
        #          detail="Request must use only one data source: either query parameters or body, not both",
        #      )

        # -------------------------
        # 6. JSON body inspection
        # -------------------------
        if content_type.startswith("application/json"):
            body = await request.body()

            if len(body) > MAX_JSON_BODY_SIZE:
                raise HTTPException(
                    status_code=422,
                    detail="JSON body too large",
                )

            if body:
                try:
                    payload = json.loads(body)
                except json.JSONDecodeError:
                    raise HTTPException(
                        status_code=422,
                        detail="Invalid JSON body",
                    )

                inspect_json(payload)

        # -------------------------
        # 7. Content-Type sanity (Reject unknown)
        # -------------------------
        if content_type and not any(
            content_type.startswith(t)
            for t in (
                "application/json",
                "multipart/form-data",
                "application/x-www-form-urlencoded",
            )
        ) and has_body:
            raise HTTPException(
                status_code=422,
                detail="Unsupported Content-Type",
            )

        return await call_next(request)