be-optimumoh/src/middleware.py

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

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

ALLOWED_MULTI_PARAMS = {
    "sortBy[]",
    "descending[]",
    "exclude[]",
    "assetnums",
    "plant_ids",
    "job_ids",
}

ALLOWED_DATA_PARAMS = {
    "actual_shutdown", "all_params", "analysis_metadata", "asset_contributions",
    "assetnum", "assetnums", "assigned_date", "availability", "availableScopes",
    "avg_cost", "birbaum", "calculation_type", "capacity_weight", "code",
    "contribution", "corrective_cost", "corrective_costs", "cost", "costPerFailure",
    "cost_savings_vs_planned", "cost_threshold", "cost_trend", "created_at",
    "crew_number", "criticalParts", "critical_procurement_items", "current_eaf",
    "current_plant_eaf", "current_stock", "current_user", "cut_hours",
    "daily_failures", "data", "datetime", "day", "days", "descending",
    "description", "down_time", "duration", "duration_oh", "eaf_gap",
    "eaf_improvement_text", "eaf_input", "efficiency", "end_date",
    "equipment_name", "equipment_results", "equipment_with_sparepart_constraints",
    "exclude", "excluded_equipment", "expected_delivery_date", "filter_spec",
    "finish", "fleet_statistics", "id", "improvement_impact", "included_equipment",
    "included_in_optimization", "intervalDays", "is_included", "itemnum",
    "items", "itemsPerPage", "items_per_page", "job", "job_ids",
    "last_overhaul_date", "lead_time", "location", "location_tag", "location_tags",
    "maintenance_type", "master_equipment", "material_cost", "max_interval",
    "max_interval_months", "message", "month", "months_from_planned", "name",
    "next_planned_overhaul", "node", "num_failures", "num_of_failures",
    "ohSessionId", "oh_scope", "oh_session_id", "oh_type", "oh_types",
    "optimal_analysis", "optimal_breakdown", "optimal_month", "optimal_total_cost",
    "optimization_success", "optimum_analysis", "optimum_day", "optimum_oh",
    "optimum_oh_day", "optimum_oh_month", "order_date", "overhaulCost",
    "overhaul_activity", "overhaul_cost", "overhaul_costs",
    "overhaul_reference_type", "overhaul_scope", "overhaul_scope_id", "overview",
    "page", "parent", "parent_id", "plan_duration", "planned_month",
    "planned_outage", "plant_level_benefit", "po_pr_id", "po_vendor_delivery_date",
    "possible_plant_eaf", "priority_score", "procurement_cost", "procurement_costs",
    "procurement_details", "projected_eaf_improvement", "quantity",
    "quantity_required", "query_str", "recommendedScope", "recommended_reduced_outage",
    "reference", "reference_id", "remark", "removal_date", "required_improvement",
    "results", "schedules", "scope", "scope_calculation_id", "scope_equipment_job",
    "scope_name", "scope_overhaul", "service_cost", "session", "simulation",
    "simulation_id", "sort_by", "sortBy[]", "descending[]", "exclude[]",
    "sparepart_id", "sparepart_impact", "sparepart_name", "sparepart_summary",
    "spreadsheet_link", "start", "start_date", "status", "subsystem", "system",
    "systemComponents", "target_plant_eaf", "tasks", "timing_recommendation",
    "total", "totalPages", "total_cost", "total_equipment", "total_equipment_analyzed",
    "total_procurement_items", "type", "unit_cost", "warning_message", "with_results",
    "workscope", "workscope_group", "year", "_", "t", "timestamp",
    "q", "filter", "currentUser", "risk_cost", "all", "with_results",
    "eaf_threshold", "simulation_id", "scope_calculation_id", "calculation_id"
}

ALLOWED_HEADERS = {
    "host",
    "user-agent",
    "accept",
    "accept-language",
    "accept-encoding",
    "connection",
    "upgrade-insecure-requests",
    "if-modified-since",
    "if-none-match",
    "cache-control",
    "authorization",
    "content-type",
    "content-length",
    "origin",
    "referer",
    "sec-fetch-dest",
    "sec-fetch-mode",
    "sec-fetch-site",
    "sec-fetch-user",
    "sec-ch-ua",
    "sec-ch-ua-mobile",
    "sec-ch-ua-platform",
    "pragma",
    "dnt",
    "priority",
    "x-forwarded-for",
    "x-forwarded-proto",
    "x-forwarded-host",
    "x-forwarded-port",
    "x-real-ip",
    "x-request-id",
    "x-correlation-id",
    "x-requested-with",
    "x-csrf-token",
    "x-xsrf-token",
    "postman-token",
    "x-forwarded-path",
    "x-forwarded-prefix",
    "cookie",
    "x-kong-request-id"
}

MAX_QUERY_PARAMS = 50
MAX_QUERY_LENGTH = 2000
MAX_JSON_BODY_SIZE = 1024 * 500  # 500 KB

XSS_PATTERN = re.compile(
    r"("
    r"<(script|iframe|embed|object|svg|img|video|audio|base|link|meta|form|button|details|animate)\b|"
    r"javascript\s*:|vbscript\s*:|data\s*:[^,]*base64[^,]*|data\s*:text/html|"
    r"\bon[a-z]+\s*=|" # Catch-all for any 'on' event (onerror, onclick, etc.)
    r"style\s*=.*expression\s*\(|" # Old IE specific
    r"\b(eval|setTimeout|setInterval|Function)\s*\("
    r")",
    re.IGNORECASE,
)

SQLI_PATTERN = re.compile(
    r"("
    # 1. Keywords followed by whitespace and common SQL characters
    r"\b(UNION|SELECT|INSERT|UPDATE|DELETE|DROP|ALTER|CREATE|TRUNCATE|EXEC(UTE)?|DECLARE)\b\s+[\w\*\(\']|"
    
    # 2. Time-based attacks (more specific than just 'SLEEP')
    r"\b(WAITFOR\b\s+DELAY|PG_SLEEP|SLEEP\s*\()|"
    
    # 3. System tables/functions
    r"\b(INFORMATION_SCHEMA|SYS\.|SYSOBJECTS|XP_CMDSHELL|LOAD_FILE|INTO\s+OUTFILE)\b|"
    
    # 4. Logical Tautologies (OR 1=1) - Optimized for boundaries
    r"\b(OR|AND)\b\s+['\"]?\d+['\"]?\s*=\s*['\"]?\d+|"
    
    # 5. Comments
    # Match '--' if at start or preceded by whitespace
    r"(?<!\S)--|"
    # Match block comments, ensuring they aren't part of mime patterns like */*
    r"(?<!\*)/\*|(?<!\*)\*/(?!\*)|"
    # Match '#' if at start or preceded by whitespace
    r"(?<!\S)#|"
    
    # 6. Hex / Stacked Queries
    r";\s*\b(SELECT|DROP|DELETE|UPDATE|INSERT)\b"
    r")",
    re.IGNORECASE
)

RCE_PATTERN = re.compile(
    r"("
    r"\$\(.*\)|`.*`|"  # Command substitution $(...) or `...`
    r"[;&|]\s*(cat|ls|id|whoami|pwd|ifconfig|ip|netstat|nc|netcat|nmap|curl|wget|python|php|perl|ruby|bash|sh|cmd|powershell|pwsh|sc\s+|tasklist|taskkill|base64|sudo|crontab|ssh|ftp|tftp)|"
    # Only flag naked commands if they are clearly standalone or system paths
    r"\b(/etc/passwd|/etc/shadow|/etc/group|/etc/issue|/proc/self/|/windows/system32/|C:\\Windows\\)\b"
    r")",
    re.IGNORECASE,
)

TRAVERSAL_PATTERN = re.compile(
    r"(\.\.[/\\]|%2e%2e%2f|%2e%2e/|\.\.%2f|%2e%2e%5c|%252e%252e%252f|\\00)",
    re.IGNORECASE,
)

FORBIDDEN_JSON_KEYS = {"__proto__", "constructor", "prototype"}

DYNAMIC_KEYS = {
    "data",
    "results",
    "analysis_metadata",
    "asset_contributions",
    "equipment_results",
    "optimal_analysis",
    "optimum_analysis",
    "schedules",
    "tasks",
    "all_params",
    "parameters",
    "program_data"
}

log = logging.getLogger("security_logger")


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 not isinstance(value, str) or value == "*/*":
        return
    

    if XSS_PATTERN.search(value):
        log.warning(f"Security violation: Potential XSS payload detected in {source}, value: {value}")
        raise HTTPException(
            status_code=422,
            detail="Invalid request parameters",
        )

    if SQLI_PATTERN.search(value):
        log.warning(f"Security violation: Potential SQL injection payload detected in {source}, value: {value}")
        raise HTTPException(
            status_code=422,
            detail="Invalid request parameters",
        )

    if RCE_PATTERN.search(value):
        log.warning(f"Security violation: Potential RCE payload detected in {source}, value: {value}")
        raise HTTPException(
            status_code=422,
            detail="Invalid request parameters",
        )

    if TRAVERSAL_PATTERN.search(value):
        log.warning(f"Security violation: Potential Path Traversal payload detected in {source}, value: {value}")
        raise HTTPException(
            status_code=422,
            detail="Invalid request parameters",
        )

    if has_control_chars(value):
        log.warning(f"Security violation: Invalid control characters detected in {source}")
        raise HTTPException(
            status_code=422,
            detail="Invalid request parameters",
        )


def inspect_json(obj, path="body", check_whitelist=True):
    if isinstance(obj, dict):
        for key, value in obj.items():
            if key in FORBIDDEN_JSON_KEYS:
                log.warning(f"Security violation: Forbidden JSON key detected: {path}.{key}")
                raise HTTPException(
                    status_code=422,
                    detail="Invalid request parameters",
                )
            
            if check_whitelist and key not in ALLOWED_DATA_PARAMS:
                log.warning(f"Security violation: Unknown JSON key detected: {path}.{key}")
                raise HTTPException(
                    status_code=422,
                    detail="Invalid request parameters",
                )
            
            # Recurse. If the key is a dynamic container, we stop whitelist checking for children.
            should_check_subkeys = check_whitelist and (key not in DYNAMIC_KEYS)
            inspect_json(value, f"{path}.{key}", check_whitelist=should_check_subkeys)
    elif isinstance(obj, list):
        for i, item in enumerate(obj):
            inspect_json(item, f"{path}[{i}]", check_whitelist=check_whitelist)
    elif isinstance(obj, str):
        inspect_value(obj, path)


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

class RequestValidationMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next):
        # -------------------------
        # 0. Header validation
        # -------------------------
        header_keys = [key.lower() for key, _ in request.headers.items()]
        
        # Check for duplicate headers
        header_counter = Counter(header_keys)
        duplicate_headers = [key for key, count in header_counter.items() if count > 1]
        
        ALLOW_DUPLICATE_HEADERS = {'accept', 'accept-encoding', 'accept-language', 'accept-charset', 'cookie'}
        real_duplicates = [h for h in duplicate_headers if h not in ALLOW_DUPLICATE_HEADERS]
        if real_duplicates:
            log.warning(f"Security violation: Duplicate headers detected: {real_duplicates}")
            raise HTTPException(
                status_code=422,
                detail="Invalid request parameters",
            )
            
        # Whitelist headers
        unknown_headers = [key for key in header_keys if key not in ALLOWED_HEADERS]
        if unknown_headers:
             filtered_unknown = [h for h in unknown_headers if not h.startswith('sec-')]
             if filtered_unknown:
                log.warning(f"Security violation: Unknown headers detected: {filtered_unknown}")
                raise HTTPException(
                    status_code=422,
                    detail="Invalid request parameters",
                )
             
        # Inspect header values
        for key, value in request.headers.items():
            if value:
                inspect_value(value, f"header '{key}'")

        # -------------------------
        # 1. Query string limits
        # -------------------------
        if len(request.url.query) > MAX_QUERY_LENGTH:
            log.warning(f"Security violation: Query string too long")
            raise HTTPException(
                status_code=422,
                detail="Invalid request parameters",
            )

        params = request.query_params.multi_items()

        if len(params) > MAX_QUERY_PARAMS:
            log.warning(f"Security violation: Too many query parameters")
            raise HTTPException(
                status_code=422,
                detail="Invalid request parameters",
            )

        # Check for unknown query parameters
        unknown_params = [key for key, _ in params if key not in ALLOWED_DATA_PARAMS]
        if unknown_params:
            log.warning(f"Security violation: Unknown query parameters detected: {unknown_params}")
            raise HTTPException(
                status_code=422,
                detail="Invalid request parameters",
            )

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

        if duplicates:
            log.warning(f"Security violation: Duplicate query parameters detected: {duplicates}")
            raise HTTPException(
                status_code=422,
                detail="Invalid request parameters",
            )

        # -------------------------
        # 3. Query param inspection & Pagination
        # -------------------------
        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}'")
            
            if key in pagination_size_keys and value:
                try:
                    size_val = int(value)
                    if size_val > 50:
                        log.warning(f"Security violation: Pagination size too large ({size_val})")
                        raise HTTPException(
                            status_code=422,
                            detail="Invalid request parameters",
                        )
                    if size_val % 5 != 0:
                        log.warning(f"Security violation: Pagination size not multiple of 5 ({size_val})")
                        raise HTTPException(
                            status_code=422,
                            detail="Invalid request parameters",
                        )
                except ValueError:
                    log.warning(f"Security violation: Pagination size invalid value")
                    raise HTTPException(
                        status_code=422,
                        detail="Invalid request parameters",
                    )

        # -------------------------
        # 4. Content-Type sanity
        # -------------------------
        content_type = request.headers.get("content-type", "")
        if content_type and not any(
            content_type.startswith(t)
            for t in ("application/json", "multipart/form-data", "application/x-www-form-urlencoded")
        ):
            log.warning(f"Security violation: Unsupported Content-Type: {content_type}")
            raise HTTPException(status_code=422, detail="Invalid request parameters")

        # -------------------------
        # 5. Single source check (Query vs JSON Body)
        # -------------------------
        has_query = len(params) > 0
        has_body = False
        
        if content_type.startswith("application/json"):
            # We can't easily check body existence without consuming it, 
            # so we check if Content-Length > 0
            content_length = request.headers.get("content-length")
            if content_length and int(content_length) > 0:
                has_body = True
                
        if has_query and has_body:
            log.warning(f"Security violation: Mixed parameters (query + JSON body)")
            raise HTTPException(
                status_code=422,
                detail="Invalid request parameters",
            )

        # -------------------------
        # 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:
                    log.warning(f"Security violation: Invalid JSON body")
                    raise HTTPException(status_code=422, detail="Invalid request parameters")

                inspect_json(payload)

            # Re-inject body for downstream handlers
            async def receive():
                return {"type": "http.request", "body": body}
            request._receive = receive

        return await call_next(request)