Prowler Mcp

# Example: BaseTool Abstract Class
# Source: mcp_server/prowler_mcp_server/prowler_app/tools/base.py

import inspect
from abc import ABC
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from fastmcp import FastMCP

from prowler_mcp_server.lib.logger import logger
from prowler_mcp_server.prowler_app.utils.api_client import ProwlerAPIClient


class BaseTool(ABC):
    """
    Abstract base class for MCP tools.

    Key patterns:
    1. Auto-registers all public async methods as tools
    2. Provides shared api_client and logger via properties
    3. Subclasses just define async methods with Field() parameters
    """

    def __init__(self):
        self._api_client = ProwlerAPIClient()
        self._logger = logger

    @property
    def api_client(self) -> ProwlerAPIClient:
        """Shared API client for making authenticated requests."""
        return self._api_client

    @property
    def logger(self):
        """Logger for structured logging."""
        return self._logger

    def register_tools(self, mcp: "FastMCP") -> None:
        """
        Auto-register all public async methods as MCP tools.

        Subclasses don't need to override this - just define async methods.
        """
        registered_count = 0

        for name, method in inspect.getmembers(self, predicate=inspect.ismethod):
            # Skip private/protected methods
            if name.startswith("_"):
                continue
            # Skip inherited methods
            if name in ["register_tools", "api_client", "logger"]:
                continue
            # Only register async methods
            if inspect.iscoroutinefunction(method):
                mcp.tool(method)
                registered_count += 1
                self.logger.debug(f"Auto-registered tool: {name}")

        self.logger.info(
            f"Auto-registered {registered_count} tools from {self.__class__.__name__}"
        )

# Example: MCP Models with MinimalSerializerMixin
# Source: mcp_server/prowler_mcp_server/prowler_app/models/

from typing import Any, Literal

from pydantic import (
    BaseModel,
    ConfigDict,
    Field,
    SerializerFunctionWrapHandler,
    model_serializer,
)


class MinimalSerializerMixin(BaseModel):
    """
    Mixin that excludes empty values from serialization.

    Key pattern: Reduces token usage by removing None, empty strings, empty lists/dicts.
    Use this for all LLM-facing models.
    """

    @model_serializer(mode="wrap")
    def _serialize(self, handler: SerializerFunctionWrapHandler) -> dict[str, Any]:
        data = handler(self)
        return {k: v for k, v in data.items() if not self._should_exclude(k, v)}

    def _should_exclude(self, key: str, value: Any) -> bool:
        """Override in subclasses for custom exclusion logic."""
        if value is None:
            return True
        if value == "":
            return True
        if isinstance(value, list) and not value:
            return True
        if isinstance(value, dict) and not value:
            return True
        return False


class CheckRemediation(MinimalSerializerMixin, BaseModel):
    """Remediation information - uses mixin to strip empty fields."""

    model_config = ConfigDict(frozen=True)

    cli: str | None = Field(default=None, description="CLI command for remediation")
    terraform: str | None = Field(default=None, description="Terraform code")
    other: str | None = Field(default=None, description="Other remediation steps")
    recommendation: str | None = Field(
        default=None, description="Best practice recommendation"
    )


class SimplifiedFinding(MinimalSerializerMixin, BaseModel):
    """
    Lightweight finding for list responses.

    Key pattern: Two-tier serialization
    - SimplifiedFinding: minimal fields for lists (fast, low tokens)
    - DetailedFinding: full fields for single item (complete info)
    """

    model_config = ConfigDict(frozen=True)

    id: str = Field(description="Finding UUID")
    uid: str = Field(description="Unique finding identifier")
    status: Literal["FAIL", "PASS", "MANUAL"] = Field(description="Finding status")
    severity: str = Field(description="Severity level")
    check_id: str = Field(description="Check ID that generated this finding")
    resource_name: str | None = Field(default=None, description="Affected resource")

    @classmethod
    def from_api_response(cls, data: dict) -> "SimplifiedFinding":
        """Transform JSON:API response to model."""
        attributes = data["attributes"]
        return cls(
            id=data["id"],
            uid=attributes["uid"],
            status=attributes["status"],
            severity=attributes["severity"],
            check_id=attributes["check_id"],
            resource_name=attributes.get("resource_name"),
        )


class DetailedFinding(SimplifiedFinding):
    """
    Full finding details - extends SimplifiedFinding.

    Key pattern: Inheritance for two-tier serialization.
    """

    status_extended: str = Field(description="Detailed status message")
    region: str | None = Field(default=None, description="Cloud region")
    remediation: CheckRemediation | None = Field(default=None, description="How to fix")

    @classmethod
    def from_api_response(cls, data: dict) -> "DetailedFinding":
        """Transform JSON:API response to detailed model."""
        attributes = data["attributes"]
        check_metadata = attributes.get("check_metadata", {})
        remediation_data = check_metadata.get("Remediation", {})

        return cls(
            id=data["id"],
            uid=attributes["uid"],
            status=attributes["status"],
            severity=attributes["severity"],
            check_id=attributes["check_id"],
            resource_name=attributes.get("resource_name"),
            status_extended=attributes.get("status_extended", ""),
            region=attributes.get("region"),
            remediation=(
                CheckRemediation(
                    cli=remediation_data.get("Code", {}).get("CLI"),
                    terraform=remediation_data.get("Code", {}).get("Terraform"),
                    recommendation=remediation_data.get("Recommendation", {}).get(
                        "Text"
                    ),
                )
                if remediation_data
                else None
            ),
        )


class FindingsListResponse(BaseModel):
    """Wrapper for list responses with pagination."""

    findings: list[SimplifiedFinding]
    total: int
    page: int
    page_size: int

    @classmethod
    def from_api_response(cls, data: dict) -> "FindingsListResponse":
        findings = [
            SimplifiedFinding.from_api_response(f) for f in data.get("data", [])
        ]
        meta = data.get("meta", {}).get("pagination", {})
        return cls(
            findings=findings,
            total=meta.get("count", len(findings)),
            page=meta.get("page", 1),
            page_size=meta.get("page_size", len(findings)),
        )

# Example: Tool Implementation (FindingsTools)
# Source: mcp_server/prowler_mcp_server/prowler_app/tools/findings.py

from typing import Any, Literal

from prowler_mcp_server.prowler_app.models.findings import (
    DetailedFinding,
    FindingsListResponse,
)
from prowler_mcp_server.prowler_app.tools.base import BaseTool
from pydantic import Field


class FindingsTools(BaseTool):
    """
    MCP tools for security findings.

    Key patterns:
    1. Extends BaseTool (no need to override register_tools)
    2. Each async method becomes a tool automatically
    3. Use pydantic.Field() for parameter documentation
    4. Return dict from model_dump() for serialization
    """

    async def search_security_findings(
        self,
        severity: list[
            Literal["critical", "high", "medium", "low", "informational"]
        ] = Field(
            default=[],
            description="Filter by severity levels. Multiple values allowed.",
        ),
        status: list[Literal["FAIL", "PASS", "MANUAL"]] = Field(
            default=["FAIL"],
            description="Filter by finding status. Default: ['FAIL'].",
        ),
        provider_type: list[str] = Field(
            default=[],
            description="Filter by cloud provider (aws, azure, gcp, etc.).",
        ),
        page_size: int = Field(
            default=50,
            description="Number of results per page.",
        ),
        page_number: int = Field(
            default=1,
            description="Page number (1-indexed).",
        ),
    ) -> dict[str, Any]:
        """
        Search security findings with rich filtering.

        Returns simplified finding data optimized for LLM consumption.
        """
        # Validate page size
        self.api_client.validate_page_size(page_size)

        # Build query parameters
        params = {
            "page[size]": page_size,
            "page[number]": page_number,
        }
        if severity:
            params["filter[severity__in]"] = ",".join(severity)
        if status:
            params["filter[status__in]"] = ",".join(status)
        if provider_type:
            params["filter[provider_type__in]"] = ",".join(provider_type)

        # Make API request
        api_response = await self.api_client.get("/findings", params=params)

        # Transform to simplified model and return
        simplified_response = FindingsListResponse.from_api_response(api_response)
        return simplified_response.model_dump()

    async def get_finding_details(
        self,
        finding_id: str = Field(
            description="UUID of the finding to retrieve.",
        ),
    ) -> dict[str, Any]:
        """
        Get comprehensive details for a specific finding.

        Returns full finding data including remediation steps.
        """
        params = {"include": "resources,scan"}
        api_response = await self.api_client.get(
            f"/findings/{finding_id}", params=params
        )
        detailed_finding = DetailedFinding.from_api_response(
            api_response.get("data", {})
        )
        return detailed_finding.model_dump()