prowler/mcp_server/prowler_mcp_server/prowler_app/utils/auth.py

import base64
import json
import os
from datetime import datetime

from fastmcp.server.dependencies import get_http_headers

from prowler_mcp_server import __version__
from prowler_mcp_server.lib.errors import CredentialError
from prowler_mcp_server.lib.logger import logger


class ProwlerAppAuth:
    """Handles authentication for Prowler API using API keys or JWT tokens."""

    def __init__(
        self,
        mode: str = os.getenv("PROWLER_MCP_TRANSPORT_MODE", "stdio"),
        base_url: str = os.getenv("API_BASE_URL", "https://api.prowler.com/api/v1"),
    ):
        self.base_url = base_url.rstrip("/")
        logger.info(f"Using Prowler API base URL: {self.base_url}")
        self.mode = mode
        self.access_token: str | None = None
        self.api_key: str | None = None

        if mode == "stdio":  # STDIO mode
            # PROWLER_API_KEY is the current variable; PROWLER_APP_API_KEY is kept
            # as a backward-compatible fallback so existing setups keep working.
            self.api_key = os.getenv("PROWLER_API_KEY") or os.getenv(
                "PROWLER_APP_API_KEY"
            )

            if not self.api_key:
                raise ValueError("PROWLER_API_KEY environment variable is required")

            if not self.api_key.startswith("pk_"):
                raise ValueError("Prowler API key format is incorrect")

    def _parse_jwt(self, token: str) -> dict | None:
        """Parse JWT token and return payload

        Args:
            token: JWT token to parse

        Returns:
            Parsed JWT payload, or None if parsing fails
        """
        if not token:
            return None

        try:
            parts = token.split(".")
            if len(parts) != 3:
                return None

            # Decode base64url
            base64_payload = parts[1]
            # Replace base64url characters
            base64_payload = base64_payload.replace("-", "+").replace("_", "/")

            # Add padding if necessary
            while len(base64_payload) % 4:
                base64_payload += "="

            # Decode and parse JSON
            decoded = base64.b64decode(base64_payload).decode("utf-8")
            payload = json.loads(decoded)

            # A JWT payload is a JSON object. A list or a scalar decodes just as
            # cleanly, so the type is checked here rather than left to blow up as
            # an AttributeError on the first claim read.
            return payload if isinstance(payload, dict) else None
        except Exception as e:
            logger.warning(f"Failed to parse JWT token: {e}")
            return None

    async def authenticate(self) -> str:
        """Authenticate and return token (API key for STDIO, API key or JWT for HTTP)."""
        if self.mode == "http":
            headers = get_http_headers(include={"authorization"})
            authorization_header = headers.get("authorization", None)

            if not authorization_header:
                raise CredentialError("No Authorization header was sent")

            # Extract token from Bearer header. Authentication scheme names are
            # case-insensitive (RFC 7235), and only the scheme prefix is removed:
            # a token that happens to contain the word again keeps it.
            scheme, _, credential = authorization_header.partition(" ")
            token = credential.strip()
            if scheme.lower() != "bearer" or not token:
                raise CredentialError(
                    "The Authorization header is not in 'Bearer <token>' form"
                )

            # Check if it's an API key or JWT token
            if token.startswith("pk_"):
                # API key - no expiration check needed
                return token
            else:
                # JWT token - validate and check expiration
                payload = self._parse_jwt(token)
                if not payload:
                    raise CredentialError("The token is not a readable JWT")

                # Check if token is expired. `exp` is a numeric date in the
                # spec, so a missing or non-numeric one makes the token
                # unusable rather than merely stale -- comparing it would raise
                # a TypeError and leave the failure masked as unclassified.
                exp = payload.get("exp")
                if isinstance(exp, bool) or not isinstance(exp, (int, float)):
                    raise CredentialError(
                        "The token carries no readable 'exp' expiration claim"
                    )

                now = int(datetime.now().timestamp())
                if exp <= now:
                    raise CredentialError("The token has expired")

                return token
        else:
            # PROWLER_MCP_TRANSPORT_MODE holds something this server does not
            # support. Nothing about a call caused it and nothing about a call
            # can fix it, so it stays unclassified: masked for the model, logged
            # for whoever runs the server.
            raise RuntimeError(f"Invalid mode: {self.mode}")

    async def get_valid_token(self) -> str:
        """Get a valid token (API key or JWT token)."""
        if self.mode == "stdio" and self.api_key:
            return self.api_key
        else:
            return await self.authenticate()

    def get_headers(self, token: str) -> dict[str, str]:
        """Get headers for API requests with authentication."""
        if token.startswith("pk_"):
            authorization_header = f"Api-Key {token}"
        else:
            authorization_header = f"Bearer {token}"

        headers = {
            "Authorization": authorization_header,
            "Content-Type": "application/vnd.api+json",
            "Accept": "application/vnd.api+json",
            "User-Agent": f"prowler-mcp-server/{__version__}",
        }

        return headers