Extending the MCP Server

mcp_server/prowler_mcp_server/
├── server.py                 # Main orchestrator
├── main.py                   # CLI entry point
├── prowler_hub/
├── prowler_app/
│   ├── tools/                # Tool implementations
│   ├── models/               # Pydantic models
│   └── utils/                # API client, auth, loader
└── prowler_documentation/

# prowler_app/tools/new_feature.py
from typing import Any

from pydantic import Field

from prowler_mcp_server.prowler_app.models.new_feature import (
    FeatureListResponse,
    DetailedFeature,
)
from prowler_mcp_server.prowler_app.tools.base import BaseTool


class NewFeatureTools(BaseTool):
    """Tools for managing new features."""

    async def list_features(
        self,
        status: str | None = Field(
            default=None,
            description="Filter by status (active, inactive, pending)"
        ),
        page_size: int = Field(
            default=50,
            description="Number of results per page (1-100)"
        ),
    ) -> dict[str, Any]:
        """List all features with optional filtering.

        Returns a lightweight list of features optimized for LLM consumption.
        Use get_feature for complete information about a specific feature.
        """
        # Validate parameters
        self.api_client.validate_page_size(page_size)

        # Build query parameters
        params: dict[str, Any] = {"page[size]": page_size}
        if status:
            params["filter[status]"] = status

        # Make API request
        clean_params = self.api_client.build_filter_params(params)
        response = await self.api_client.get("/api/v1/features", params=clean_params)

        # Transform to LLM-friendly format
        return FeatureListResponse.from_api_response(response).model_dump()

    async def get_feature(
        self,
        feature_id: str = Field(description="The UUID of the feature"),
    ) -> dict[str, Any]:
        """Get detailed information about a specific feature.

        Returns complete feature details including configuration and metadata.
        """
        try:
            response = await self.api_client.get(f"/api/v1/features/{feature_id}")
            return DetailedFeature.from_api_response(response["data"]).model_dump()
        except Exception as e:
            self.logger.error(f"Failed to get feature {feature_id}: {e}")
            return {"error": str(e), "status": "failed"}

# prowler_app/models/new_feature.py
from typing import Any

from pydantic import Field

from prowler_mcp_server.prowler_app.models.base import MinimalSerializerMixin


class SimplifiedFeature(MinimalSerializerMixin):
    """Lightweight feature for list operations."""

    id: str = Field(description="Unique feature identifier")
    name: str = Field(description="Feature name")
    status: str = Field(description="Current status")

    @classmethod
    def from_api_response(cls, data: dict[str, Any]) -> "SimplifiedFeature":
        """Transform API response to simplified format."""
        attributes = data.get("attributes", {})
        return cls(
            id=data["id"],
            name=attributes["name"],
            status=attributes["status"],
        )


class DetailedFeature(SimplifiedFeature):
    """Extended feature with complete details."""

    description: str | None = Field(default=None, description="Feature description")
    configuration: dict[str, Any] | None = Field(default=None, description="Configuration")
    created_at: str = Field(description="Creation timestamp")
    updated_at: str = Field(description="Last update timestamp")

    @classmethod
    def from_api_response(cls, data: dict[str, Any]) -> "DetailedFeature":
        """Transform API response to detailed format."""
        attributes = data.get("attributes", {})
        return cls(
            id=data["id"],
            name=attributes["name"],
            status=attributes["status"],
            description=attributes.get("description"),
            configuration=attributes.get("configuration"),
            created_at=attributes["created_at"],
            updated_at=attributes["updated_at"],
        )


class FeatureListResponse(MinimalSerializerMixin):
    """Response wrapper for feature list operations."""

    count: int = Field(description="Total number of features")
    features: list[SimplifiedFeature] = Field(description="List of features")

    @classmethod
    def from_api_response(cls, response: dict[str, Any]) -> "FeatureListResponse":
        """Transform API response to list format."""
        data = response.get("data", [])
        features = [SimplifiedFeature.from_api_response(item) for item in data]
        return cls(count=len(features), features=features)

INFO - Auto-registered 2 tools from NewFeatureTools
INFO - Loaded and registered: NewFeatureTools

# prowler_hub/server.py
from fastmcp import FastMCP

hub_mcp_server = FastMCP("prowler-hub")

@hub_mcp_server.tool()
async def get_new_artifact(
    artifact_id: str,
) -> dict:
    """Fetch a specific artifact from Prowler Hub.

    Args:
        artifact_id: The unique identifier of the artifact

    Returns:
        Dictionary containing artifact details
    """
    response = prowler_hub_client.get(f"/artifact/{artifact_id}")
    response.raise_for_status()
    return response.json()

from prowler_mcp_server.prowler_app.models.base import MinimalSerializerMixin

class MyModel(MinimalSerializerMixin):
    """Model that excludes empty values from serialization."""
    required_field: str
    optional_field: str | None = None  # Excluded if None
    empty_list: list = []              # Excluded if empty

class SimplifiedItem(MinimalSerializerMixin):
    """Use for list operations - minimal fields."""
    id: str
    name: str
    status: str

class DetailedItem(SimplifiedItem):
    """Use for get operations - extends simplified with details."""
    description: str | None = None
    configuration: dict | None = None
    created_at: str
    updated_at: str

@classmethod
def from_api_response(cls, data: dict[str, Any]) -> "MyModel":
    """Transform API response to model.

    This method handles the JSON:API format used by Prowler API,
    extracting attributes and relationships as needed.
    """
    attributes = data.get("attributes", {})
    return cls(
        id=data["id"],
        name=attributes["name"],
        # ... map other fields
    )

class MyTools(BaseTool):
    async def my_tool(self) -> dict:
        # GET request
        response = await self.api_client.get("/api/v1/endpoint", params={"key": "value"})

        # POST request
        response = await self.api_client.post(
            "/api/v1/endpoint",
            json_data={"data": {"type": "items", "attributes": {...}}}
        )

        # PATCH request
        response = await self.api_client.patch(
            f"/api/v1/endpoint/{id}",
            json_data={"data": {"attributes": {...}}}
        )

        # DELETE request
        response = await self.api_client.delete(f"/api/v1/endpoint/{id}")

# Validate page size (1-1000)
self.api_client.validate_page_size(page_size)

# Normalize date range with max days limit
date_range = self.api_client.normalize_date_range(date_from, date_to, max_days=2)

# Build filter parameters (handles type conversion)
clean_params = self.api_client.build_filter_params({
    "filter[status]": "active",
    "filter[severity__in]": ["high", "critical"],  # Converts to comma-separated
    "filter[muted]": True,  # Converts to "true"
})

# Poll async task until completion
result = await self.api_client.poll_task_until_complete(
    task_id=task_id,
    timeout=60,
    poll_interval=1.0
)

async def search_items(self, status: str = Field(...)) -> dict:
    """Search items with advanced filtering.

    Returns a lightweight list optimized for LLM consumption.
    Use get_item for complete details about a specific item.

    Common workflows:
    - Find critical items: status="critical"
    - Find recent items: Use date_from parameter
    """

async def get_item(self, item_id: str) -> dict:
    try:
        response = await self.api_client.get(f"/api/v1/items/{item_id}")
        return DetailedItem.from_api_response(response["data"]).model_dump()
    except Exception as e:
        self.logger.error(f"Failed to get item {item_id}: {e}")
        return {"error": str(e), "status": "failed"}

async def list_items(
    self,
    severity: list[str] = Field(
        default=[],
        description="Filter by severity levels (critical, high, medium, low)"
    ),
    status: str | None = Field(
        default=None,
        description="Filter by status (PASS, FAIL, MANUAL)"
    ),
    page_size: int = Field(
        default=50,
        description="Results per page"
    ),
) -> dict:

# Navigate to MCP server directory
cd mcp_server

# Run in STDIO mode (default)
uv run prowler-mcp

# Run in HTTP mode
uv run prowler-mcp --transport http --host 0.0.0.0 --port 8000

# Run with environment variables
PROWLER_API_KEY="pk_xxx" uv run prowler-mcp

cd mcp_server

uv run pytest                              # Whole suite
uv run pytest tests/prowler_app/models     # One area
uv run pytest --cov=./prowler_mcp_server   # With coverage

make test-mcp   # Runs the MCP suite exactly as CI does

from fastmcp import Client

from tests.helpers.jsonapi import jsonapi_collection, jsonapi_resource

FINDING_ATTRIBUTES = {
    "uid": "prowler-aws-s3_bucket_public_access-123456789012-us-east-1-my-bucket",
    "status": "FAIL",
    "severity": "high",
    "status_extended": "S3 bucket my-bucket is publicly accessible.",
    "delta": "new",
    "muted": False,
    "muted_reason": None,
    "check_metadata": {"checkid": "s3_bucket_public_access"},
}


async def test_search_without_dates_queries_the_latest_scan_endpoint(
    mcp_root_server, mock_api_client, mock_router
):
    """With no date range the tool targets the cheaper `/findings/latest`."""
    mock_router.add(
        "GET",
        "/api/v1/findings/latest",
        json=jsonapi_collection(
            [jsonapi_resource("findings", "f1", FINDING_ATTRIBUTES)]
        ),
    )

    async with Client(mcp_root_server) as client:
        result = await client.call_tool("prowler_search_security_findings", {})

    assert result.data["findings"][0]["check_id"] == "s3_bucket_public_access"
    assert mock_router.paths() == ["GET /api/v1/findings/latest"]