security_gateway/server.py

from __future__ import annotations

import os
import time
import logging
from enum import Enum
from typing import Any, Dict, Optional, List

import httpx
from fastmcp import FastMCP
from pydantic import BaseModel, Field
from starlette.requests import Request
from starlette.responses import JSONResponse

from .downstream_autodiscovery import DISCOVERED_TOOLS, discover_downstream_tools
from .downstream import call_downstream_tool
from .risk_model import compute_risk, risk_result_to_dict, initialize_plugins, get_plugin_cache_status
from .policy import decide_from_risk, decision_to_dict
from .sanitizer import sanitize_arguments, sanitize_output
from .audit import write_audit
from .rate_limiter import RateLimiter
from .config import RATE_LIMIT_MAX_CALLS, RATE_LIMIT_WINDOW_SECONDS

# Setup logging
logger = logging.getLogger(__name__)

# Initialize global rate limiter
_rate_limiter = RateLimiter(max_calls=RATE_LIMIT_MAX_CALLS, window_seconds=RATE_LIMIT_WINDOW_SECONDS)


class ErrorCategory(str, Enum):
    """Categorization of error types for structured error handling."""
    VALIDATION = "validation"
    NETWORK = "network"
    TIMEOUT = "timeout"
    POLICY_DENIED = "policy_denied"
    DOWNSTREAM_ERROR = "downstream_error"
    RATE_LIMITED = "rate_limited"
    UNKNOWN = "unknown"


def _normalize_downstream_result(raw: Any, success: bool, error: str | None = None) -> Dict[str, Any]:
    """Normalize downstream response to standard envelope."""
    return {
        "success": success,
        "data": raw if success else None,
        "error": error
    }


def _categorize_error(exc: Exception) -> ErrorCategory:
    """Map exception to error category."""
    msg = str(exc).lower()
    if "timeout" in msg or "timed out" in msg:
        return ErrorCategory.TIMEOUT
    if "connection" in msg or "unreachable" in msg or "refused" in msg:
        return ErrorCategory.NETWORK
    if "validation" in msg or "invalid" in msg or "schema" in msg:
        return ErrorCategory.VALIDATION
    if "downstream" in msg or "call failed" in msg:
        return ErrorCategory.DOWNSTREAM_ERROR
    return ErrorCategory.UNKNOWN


# Keys that, if sanitized/changed, indicate a potentially dangerous file operation.
CRITICAL_PATH_KEYS = {"path", "filepath", "target"}


def _critical_arg_changed(before: Any, after: Any) -> bool:
    """Recursively detect if any critical path-like argument changed due to sanitization.



    We only consider keys in CRITICAL_PATH_KEYS to avoid blocking benign PII redactions.

    """
    def walk(bv: Any, av: Any) -> bool:
        if isinstance(bv, dict) and isinstance(av, dict):
            # Direct key comparison first
            for k in CRITICAL_PATH_KEYS:
                if k in bv and k in av and bv[k] != av[k]:
                    return True
            # Recurse through all keys
            for k in set(bv.keys()) | set(av.keys()):
                if walk(bv.get(k), av.get(k)):
                    return True
            return False
        if isinstance(bv, list) and isinstance(av, list):
            for i in range(min(len(bv), len(av))):
                if walk(bv[i], av[i]):
                    return True
            return False
        return False

    try:
        return walk(before, after)
    except Exception:
        return False


def _sanitized_critical(before: Dict[str, Any], after: Dict[str, Any]) -> bool:
    """Detect if any critical key was sanitized (redacted token) or changed.



    Prefer detecting explicit redaction tokens; fall back to value inequality.

    """
    try:
        for k in CRITICAL_PATH_KEYS:
            if k in before or k in after:
                bv = before.get(k) if isinstance(before, dict) else None
                av = after.get(k) if isinstance(after, dict) else None
                if isinstance(av, str) and av.startswith("[REDACTED_"):
                    return True
                if bv is not None and av is not None and bv != av:
                    return True
        # Fallback to recursive check
        return _critical_arg_changed(before, after)
    except Exception:
        return False


mcp = FastMCP(
    name="SecurityGateway",
    instructions=(
        "A unified, secure MCP tool gateway. Exposes all tools through the "
        "'secure_call' API. Performs risk scoring, sanitization, and auditing. "
        "Use list_available_tools() to discover capabilities; always call "
        "secure_call() instead of downstream tools directly."
    ),
)


class SecureCallInput(BaseModel):
    user_id: str = Field(description="Logical user id (e.g., 'admin', 'judge-1').")
    server: str = Field(description="Downstream server key, e.g. 'filesystem'.")
    tool: str = Field(description="Tool name on the downstream server.")
    arguments: Dict[str, Any] = Field(
        default_factory=dict,
        description="Arguments to pass to the downstream tool.",
    )
    llm_context: Optional[str] = Field(
        default=None,
        description="Optional prompt or reasoning context for risk analysis.",
    )


class SecureCallOutput(BaseModel):
    allowed: bool
    redacted: bool
    reason: str
    risk_score: float
    risk_factors: List[str] = Field(default_factory=list, description="Explicit risk reasons")
    error_category: Optional[ErrorCategory] = Field(default=None, description="Categorized error type")
    execution_time_ms: float = Field(default=0.0, description="Execution time in milliseconds")
    policy: Dict[str, Any]
    risk: Dict[str, Any]
    downstream_result: Dict[str, Any] = Field(description="Normalized downstream response")
    original_arguments: Dict[str, Any] = Field(default_factory=dict, description="Raw unsanitized arguments provided")
    sanitized_arguments: Dict[str, Any] = Field(default_factory=dict, description="Arguments after sanitization")
    policy_decision: str = Field(description="High-level policy decision: allow | redacted | blocked")






async def _secure_call_impl(data: SecureCallInput) -> SecureCallOutput:
    """

    Secure proxy for ANY downstream MCP tool.



    The LLM must always use this tool instead of calling downstream servers

    directly.

    """
    start_time = time.perf_counter()
    
    # Check rate limit first
    allowed, rate_info = _rate_limiter.check_rate_limit(data.user_id)
    if not allowed:
        elapsed = (time.perf_counter() - start_time) * 1000
        error_msg = f"Rate limit exceeded: {rate_info['current_count']}/{rate_info['limit']} calls in {rate_info['window_seconds']}s. Try again in {rate_info['reset_in_seconds']}s."
        return SecureCallOutput(
            allowed=False,
            redacted=False,
            reason=error_msg,
            risk_score=0.0,
            risk_factors=["Rate limit exceeded"],
            error_category=ErrorCategory.RATE_LIMITED,
            execution_time_ms=elapsed,
            policy={"allow": False, "reason": error_msg, "rate_limit": rate_info},
            risk={"score": 0.0, "reasons": [], "flags": {}},
            downstream_result=_normalize_downstream_result(None, False, error_msg),
            original_arguments=data.arguments,
            sanitized_arguments={},
            policy_decision="blocked",
        )
    
    server_key = data.server
    tool_name = data.tool
    raw_args = data.arguments

    # Validate server is configured (even if discovery hasn't been completed yet)
    if server_key not in DISCOVERED_TOOLS:
        # Allow configured servers even if discovery hasn't completed
        from .config import DOWNSTREAM_SERVERS
        if server_key not in DOWNSTREAM_SERVERS:
            raise Exception(
                f"Unknown server '{server_key}'. "
                "Call list_available_tools() to see valid options."
            )
        # Server is configured but not yet discovered - allow it to proceed
        logger.warning(f"⚠️  Server '{server_key}' is configured but tools not discovered yet. Proceeding with request.")

    elif tool_name not in DISCOVERED_TOOLS[server_key]:
        # If discovery happened, validate the tool exists
        raise Exception(
            f"Unknown tool '{tool_name}' for server '{server_key}'. "
            "Call list_available_tools() to see valid tools."
        )

    # Risk scoring
    risk = compute_risk(
        user_id=data.user_id,
        server_key=server_key,
        tool=tool_name,
        arguments=raw_args,
        llm_context=data.llm_context,
    )
    policy = decide_from_risk(risk)

    # Always sanitize arguments before sending downstream
    sanitized_args = sanitize_arguments(raw_args)

    # If sanitization changed a critical argument (e.g., path traversal), block at gateway
    if _sanitized_critical(raw_args, sanitized_args):
        elapsed = (time.perf_counter() - start_time) * 1000 if 'start_time' in locals() else 0.0
        reason = "Blocked: sanitized critical argument (e.g., path traversal or sensitive path)."
        outcome = {"success": False, "error": reason, "category": ErrorCategory.POLICY_DENIED.value}
        write_audit(
            user_id=data.user_id,
            server_key=server_key,
            tool=tool_name,
            raw_arguments=raw_args,
            sanitized_arguments=sanitized_args,
            policy=decision_to_dict(policy),
            risk=risk_result_to_dict(risk),
            outcome=outcome,
            execution_time_ms=elapsed,
        )
        return SecureCallOutput(
            allowed=False,
            redacted=True,
            reason=reason,
            risk_score=policy.risk_score,
            risk_factors=risk.reasons,
            error_category=ErrorCategory.POLICY_DENIED,
            execution_time_ms=elapsed,
            policy=decision_to_dict(policy),
            risk=risk_result_to_dict(risk),
            downstream_result=_normalize_downstream_result(None, False, reason),
            original_arguments=raw_args,
            sanitized_arguments=sanitized_args,
            policy_decision="blocked",
        )

    if not policy.allow:
        elapsed = (time.perf_counter() - start_time) * 1000
        outcome = {"success": False, "error": "Blocked by security policy.", "category": "policy_denied"}
        write_audit(
            user_id=data.user_id,
            server_key=server_key,
            tool=tool_name,
            raw_arguments=raw_args,
            sanitized_arguments=sanitized_args,
            policy=decision_to_dict(policy),
            risk=risk_result_to_dict(risk),
            outcome=outcome,
            execution_time_ms=elapsed,
        )

        return SecureCallOutput(
            allowed=False,
            redacted=True,
            reason=policy.reason,
            risk_score=policy.risk_score,
            risk_factors=risk.reasons,
            error_category=ErrorCategory.POLICY_DENIED,
            execution_time_ms=elapsed,
            policy=decision_to_dict(policy),
            risk=risk_result_to_dict(risk),
            downstream_result=_normalize_downstream_result(None, False, "Blocked by policy"),
            original_arguments=raw_args,
            sanitized_arguments=sanitized_args,
            policy_decision="blocked",
        )

    # Allowed calls: forward to downstream MCP
    error_category = None
    try:
        downstream_raw = await call_downstream_tool(
            server_key=server_key,
            tool=tool_name,
            arguments=sanitized_args,
        )
        downstream_result = (
            sanitize_output(downstream_raw)
            if policy.redact_output
            else downstream_raw
        )
        normalized = _normalize_downstream_result(downstream_result, True)
        outcome = {"success": True, "redacted": policy.redact_output}
    except Exception as exc:
        error_category = _categorize_error(exc)
        normalized = _normalize_downstream_result(None, False, str(exc))
        outcome = {"success": False, "error": str(exc), "category": error_category.value}

    elapsed = (time.perf_counter() - start_time) * 1000
    
    # Audit everything
    write_audit(
        user_id=data.user_id,
        server_key=server_key,
        tool=tool_name,
        raw_arguments=raw_args,
        sanitized_arguments=sanitized_args,
        policy=decision_to_dict(policy),
        risk=risk_result_to_dict(risk),
        outcome=outcome,
        execution_time_ms=elapsed,
    )

    # Determine policy_decision label based on outcome
    # Distinguish between security blocks vs execution failures vs successful calls
    if error_category:
        # Execution failed - use error-specific labels
        if error_category == ErrorCategory.TIMEOUT:
            policy_decision = "timeout"
        else:
            policy_decision = "error"
    elif policy.allow:
        # Security allowed - check if output was redacted
        policy_decision = "redacted" if policy.redact_output else "allowed"
    else:
        # Security blocked the request
        policy_decision = "blocked"

    return SecureCallOutput(
        allowed=policy.allow,
        redacted=policy.redact_output,
        reason=policy.reason,
        risk_score=policy.risk_score,
        risk_factors=risk.reasons,
        error_category=error_category,
        execution_time_ms=elapsed,
        policy=decision_to_dict(policy),
        risk=risk_result_to_dict(risk),
        downstream_result=normalized,
        original_arguments=raw_args,
        sanitized_arguments=sanitized_args,
        policy_decision=policy_decision,
    )


@mcp.tool()
async def secure_call(data: SecureCallInput) -> SecureCallOutput:
    """

    MCP tool wrapper for secure downstream tool invocation.



    This is the registered MCP tool that LLMs use to safely call downstream servers.

    It delegates to _secure_call_impl for the actual implementation.

    """
    return await _secure_call_impl(data)





def _flatten_discovered_tools(tools_dict: Dict[str, Any]) -> Dict[str, Any]:
    """

    Flatten nested discovered tools into single-level namespace with server__tool_name format.



    Converts:

        {"server_key": {"tool_name": {...}}}

    To:

        {"server_key__tool_name": {...}}



    This makes tool names parsable by all LLMs (Claude, OpenAI, Gemini) consistently.

    """
    flattened = {}
    for server_key, server_tools in tools_dict.items():
        if not isinstance(server_tools, dict):
            continue
        for tool_name, tool_meta in server_tools.items():
            if tool_name.startswith("_"):
                # Skip internal metadata fields like _error
                continue
            flattened_name = f"{server_key}__{tool_name}"
            flattened[flattened_name] = tool_meta
    return flattened


@mcp.tool()
async def list_available_tools(refresh: bool = False) -> Dict[str, Any]:
        """

        List discovered downstream servers and their tools with flattened names.



        - If ``refresh`` is true, forces a re-discovery.

        - If nothing has been discovered yet, performs discovery on-demand.



        Returns tools in flattened format: server__tool_name (e.g., "ultimate_scraper__searchEventListings")

        This format is parsable by Claude, OpenAI, Gemini, and other LLMs consistently.

        """
        if refresh or not DISCOVERED_TOOLS:
            await discover_downstream_tools()
        return _flatten_discovered_tools(DISCOVERED_TOOLS)


@mcp.tool()
async def refresh_discovery() -> Dict[str, Any]:
        """Force downstream tool re-discovery and return the updated registry with flattened names."""
        await discover_downstream_tools()
        return _flatten_discovered_tools(DISCOVERED_TOOLS)


async def _on_startup() -> None:
    """

    Called once when the MCP server starts.

    Initializes the plugin system and performs downstream tool autodiscovery.

    """
    # Initialize plugin system (caches all builtin plugins)
    print("[Server] Initializing security plugins...")
    initialize_plugins()
    cache_status = get_plugin_cache_status()
    print(f"[Server] Plugin cache: {cache_status.get('loaded_plugins', 0)} plugins loaded")

    # Perform downstream tool autodiscovery
    print("[Server] Discovering downstream tools...")
    await discover_downstream_tools()


# Convenience HTTP adapter for listing available tools
@mcp.custom_route("/tools/list", methods=["GET"])  # type: ignore
async def _http_list_tools(request: Request) -> JSONResponse:
    """List all available tools from discovered downstream servers.



    Triggers discovery on-demand if the registry is empty.

    Returns both discovered tools and native functions (web_search, code_interpreter).

    """
    tools = await list_available_tools()
    return JSONResponse({"tools": tools})


@mcp.custom_route("/tools/refresh", methods=["POST"])  # type: ignore
async def _http_refresh_tools(request: Request) -> JSONResponse:
    """Force downstream autodiscovery and return the updated registry."""
    await discover_downstream_tools()
    return JSONResponse(DISCOVERED_TOOLS)


# Convenience HTTP adapter for calling native functions (web_search, code_interpreter)

@mcp.custom_route("/config/servers", methods=["GET"])  # type: ignore
async def _http_get_config(request: Request) -> JSONResponse:
    """

    Expose servers configuration for distributed clients.



    Used by Real LLM and other clients to fetch tool parameter definitions.

    Returns the entire servers.yaml configuration as JSON.

    """
    try:
        from pathlib import Path
        import yaml

        # Try multiple config paths
        config_paths = [
            Path(__file__).parent / "config" / "servers.yaml",
            Path("config/servers.yaml"),
            Path("security_gateway/config/servers.yaml"),
        ]

        config_data = None
        for config_path in config_paths:
            if config_path.exists():
                with open(config_path, 'r') as f:
                    config_data = yaml.safe_load(f)
                break

        if not config_data:
            return JSONResponse(
                {"error": "servers.yaml not found", "servers": {}},
                status_code=404
            )

        return JSONResponse(config_data)

    except Exception as e:
        return JSONResponse(
            {"error": str(e), "servers": {}},
            status_code=500
        )


@mcp.custom_route("/tools/secure_call", methods=["POST"])  # type: ignore
async def _http_secure_call(request: Request) -> JSONResponse:
    # Ensure downstream tools have been discovered before accepting calls
    if not DISCOVERED_TOOLS:
        await discover_downstream_tools()

    try:
        body = await request.json()
    except Exception:
        return JSONResponse({"error": "invalid JSON"}, status_code=400)
    
    # Extract user_id early for rate limiting
    if isinstance(body, dict) and "data" in body and isinstance(body["data"], dict):
        user_id = body["data"].get("user_id", "anonymous")
    elif isinstance(body, dict):
        user_id = body.get("user_id", "anonymous")
    else:
        user_id = "anonymous"
    
    # Check rate limit
    allowed, rate_info = _rate_limiter.check_rate_limit(user_id)
    if not allowed:
        headers = {
            "X-RateLimit-Limit": str(rate_info["limit"]),
            "X-RateLimit-Remaining": "0",
            "X-RateLimit-Reset": str(int(time.time() + rate_info["reset_in_seconds"])),
            "Retry-After": str(int(rate_info["reset_in_seconds"])),
        }
        error_msg = f"Rate limit exceeded: {rate_info['current_count']}/{rate_info['limit']} calls in {rate_info['window_seconds']}s. Try again in {rate_info['reset_in_seconds']}s."
        return JSONResponse(
            {"error": error_msg, "rate_limit": rate_info},
            status_code=429,
            headers=headers
        )

    # Accept either the raw SecureCallInput fields at top-level, or an envelope
    # like {"data": { ... }} (the MCP client sometimes uses this shape).
    if isinstance(body, dict) and "data" in body and isinstance(body["data"], dict):
        call_args = {"data": body["data"]}
    elif isinstance(body, dict):
        call_args = {"data": body}
    else:
        return JSONResponse({"error": "expected JSON object"}, status_code=400)

    try:
        # Validate and construct the input
        if "data" in call_args and isinstance(call_args["data"], dict):
            data_dict: Dict[str, Any] = call_args["data"]
        else:
            data_dict = call_args
        
        data_obj = SecureCallInput(**data_dict)
        
        # Validate server and tool have been discovered
        if data_obj.server not in DISCOVERED_TOOLS:
            raise Exception(
                f"Unknown server '{data_obj.server}'. "
                "Call list_available_tools() to see valid options."
            )

        if data_obj.tool not in DISCOVERED_TOOLS[data_obj.server]:
            raise Exception(
                f"Unknown tool '{data_obj.tool}' for server '{data_obj.server}'. "
                "Call list_available_tools() to see valid tools."
            )

        # Risk scoring
        risk = compute_risk(
            user_id=data_obj.user_id,
            server_key=data_obj.server,
            tool=data_obj.tool,
            arguments=data_obj.arguments,
            llm_context=data_obj.llm_context,
        )
        policy = decide_from_risk(risk)

        # Always sanitize arguments before sending downstream
        sanitized_args = sanitize_arguments(data_obj.arguments)

        # If sanitization changed a critical argument (e.g., path traversal), block at gateway
        if _sanitized_critical(data_obj.arguments, sanitized_args):
            elapsed = 0.0
            reason = "Blocked: sanitized critical argument (e.g., path traversal or sensitive path)."
            outcome = {"success": False, "error": reason, "category": ErrorCategory.POLICY_DENIED.value}
            write_audit(
                user_id=data_obj.user_id,
                server_key=data_obj.server,
                tool=data_obj.tool,
                raw_arguments=data_obj.arguments,
                sanitized_arguments=sanitized_args,
                policy=decision_to_dict(policy),
                risk=risk_result_to_dict(risk),
                outcome=outcome,
                execution_time_ms=elapsed,
            )

            result_obj = SecureCallOutput(
                allowed=False,
                redacted=True,
                reason=reason,
                risk_score=policy.risk_score,
                risk_factors=risk.reasons,
                error_category=ErrorCategory.POLICY_DENIED,
                execution_time_ms=elapsed,
                policy=decision_to_dict(policy),
                risk=risk_result_to_dict(risk),
                downstream_result=_normalize_downstream_result(None, False, reason),
                original_arguments=data_obj.arguments,
                sanitized_arguments=sanitized_args,
                policy_decision="blocked",
            )
            # Short-circuit HTTP handler
            try:
                payload = result_obj.model_dump()
            except Exception:
                payload = getattr(result_obj, "dict", lambda: None)() or {}
            return JSONResponse(payload)

        start_time = time.perf_counter()
        
        if not policy.allow:
            elapsed = (time.perf_counter() - start_time) * 1000
            outcome = {"success": False, "error": "Blocked by security policy.", "category": "policy_denied"}
            write_audit(
                user_id=data_obj.user_id,
                server_key=data_obj.server,
                tool=data_obj.tool,
                raw_arguments=data_obj.arguments,
                sanitized_arguments=sanitized_args,
                policy=decision_to_dict(policy),
                risk=risk_result_to_dict(risk),
                outcome=outcome,
                execution_time_ms=elapsed,
            )

            result_obj = SecureCallOutput(
                allowed=False,
                redacted=True,
                reason=policy.reason,
                risk_score=policy.risk_score,
                risk_factors=risk.reasons,
                error_category=ErrorCategory.POLICY_DENIED,
                execution_time_ms=elapsed,
                policy=decision_to_dict(policy),
                risk=risk_result_to_dict(risk),
                downstream_result=_normalize_downstream_result(None, False, "Blocked by policy"),
                original_arguments=data_obj.arguments,
                sanitized_arguments=sanitized_args,
                policy_decision="blocked",
            )
        else:
            # Allowed calls: forward to downstream MCP
            error_category = None
            try:
                downstream_raw = await call_downstream_tool(
                    server_key=data_obj.server,
                    tool=data_obj.tool,
                    arguments=sanitized_args,
                )
                downstream_result = (
                    sanitize_output(downstream_raw)
                    if policy.redact_output
                    else downstream_raw
                )
                normalized = _normalize_downstream_result(downstream_result, True)
                outcome = {"success": True, "redacted": policy.redact_output}
            except Exception as exc:
                error_category = _categorize_error(exc)
                normalized = _normalize_downstream_result(None, False, str(exc))
                outcome = {"success": False, "error": str(exc), "category": error_category.value}

            elapsed = (time.perf_counter() - start_time) * 1000
            
            # Audit everything
            write_audit(
                user_id=data_obj.user_id,
                server_key=data_obj.server,
                tool=data_obj.tool,
                raw_arguments=data_obj.arguments,
                sanitized_arguments=sanitized_args,
                policy=decision_to_dict(policy),
                risk=risk_result_to_dict(risk),
                outcome=outcome,
                execution_time_ms=elapsed,
            )

            result_obj = SecureCallOutput(
                allowed=policy.allow,
                redacted=policy.redact_output,
                reason=policy.reason,
                risk_score=policy.risk_score,
                risk_factors=risk.reasons,
                error_category=error_category,
                execution_time_ms=elapsed,
                policy=decision_to_dict(policy),
                risk=risk_result_to_dict(risk),
                downstream_result=normalized,
                original_arguments=data_obj.arguments,
                sanitized_arguments=sanitized_args,
                policy_decision="redacted" if policy.allow and policy.redact_output else ("allow" if policy.allow else "blocked"),
            )
    except Exception as exc:
        import traceback

        tb = traceback.format_exc()
        print(f"ERROR in _http_secure_call: {exc}")
        print(f"TRACEBACK:\n{tb}")
        return JSONResponse({"error": str(exc), "traceback": tb}, status_code=500)

    # `result_obj` is a Pydantic model (SecureCallOutput)
    try:
        payload = result_obj.model_dump()
    except Exception:
        # Fallback for models without model_dump
        payload = getattr(result_obj, "dict", lambda: None)() or {}
    
    # Add rate limit headers to response (check again to get current remaining count)
    _, rate_info = _rate_limiter.check_rate_limit(user_id)
    headers = {
        "X-RateLimit-Limit": str(rate_info["limit"]),
        "X-RateLimit-Remaining": str(rate_info["remaining"]),
        "X-RateLimit-Reset": str(int(time.time() + rate_info["reset_in_seconds"])),
    }

    return JSONResponse(payload, headers=headers)


@mcp.custom_route("/audit/latest", methods=["GET"])  # type: ignore
async def _http_audit_latest(request: Request) -> JSONResponse:
    """

    Retrieve the latest audit logs for the auditeye-dashboard.



    Returns the most recent audit entries from the audit.log.jsonl file.

    """
    try:
        from pathlib import Path
        import json
        from .config import AUDIT_LOG_PATH

        # Try multiple paths for the audit log (ordered by preference)
        audit_paths = [
            Path(AUDIT_LOG_PATH),  # Primary: configured path (e.g., /mnt/audit/audit.log.jsonl)
            Path(__file__).parent.parent / "audit.log.jsonl",  # Local dev
            Path("audit.log.jsonl"),  # Current working directory
            Path("./audit.log.jsonl"),  # Explicit current directory
        ]

        audit_entries = []
        audit_file = None

        # Find the audit file
        for path in audit_paths:
            if path.exists():
                audit_file = path
                break

        if not audit_file:
            return JSONResponse({"logs": [], "error": f"No audit log found at: {AUDIT_LOG_PATH}"}, status_code=200)

        # Read and parse the audit log (JSONL format)
        try:
            with open(audit_file, 'r') as f:
                for line in f:
                    if line.strip():
                        try:
                            entry = json.loads(line)
                            audit_entries.append(entry)
                        except json.JSONDecodeError:
                            continue
        except Exception as e:
            return JSONResponse(
                {"logs": [], "error": f"Failed to read audit file: {str(e)}"},
                status_code=200
            )

        # Return the last 100 entries (most recent first)
        latest = sorted(audit_entries, key=lambda x: x.get("timestamp", ""), reverse=True)[:100]

        return JSONResponse({"logs": latest})

    except Exception as e:
        import traceback
        tb = traceback.format_exc()
        return JSONResponse(
            {"logs": [], "error": str(e), "traceback": tb},
            status_code=500
        )


# FastMCP does not expose a typed 'on_startup' attribute; perform autodiscovery
# at process startup instead so tools are discovered before the server runs.
if __name__ == "__main__":
    import asyncio

    # Run discovery before starting the MCP server
    asyncio.run(discover_downstream_tools())

    # For Claude / ChatGPT: configure this as the MCP server command.
    mcp.run(transport="http", host="0.0.0.0", port=8000)