bipixie-mcp 0.3.0__py3-none-any.whl

bipixie_mcp/__init__.py

@@ -0,0 +1,57 @@
+"""
+bipixie_mcp — BI Pixie MCP Server package.
+
+Exposes the stable public surface for tests and embedders. Import-side-effect free:
+no Settings are instantiated, no credential is built, and no network I/O is performed
+at import time. This keeps stdout clean for the JSON-RPC stream (stdio transport) and
+allows tests to patch environment variables before any config object is constructed.
+"""
+
+from __future__ import annotations
+
+__version__: str = "0.3.0"
+
+# Re-export the stable public surface.  Names are imported lazily via __getattr__
+# below so that the actual modules (and their transitive imports such as fastmcp,
+# azure-identity, msal, httpx) are NOT loaded until the caller explicitly requests
+# them.  This is critical for the stdio transport: importing this package in a test
+# that patches os.environ must not trigger Settings() construction or MSAL cache
+# initialisation.
+
+__all__ = [
+    "__version__",
+    "Settings",
+    "get_settings",
+    "build_server",
+    "main",
+]
+
+# Populated on first access through __getattr__; never pre-loaded.
+_LAZY_MAP: dict[str, str] = {
+    "Settings": "bipixie_mcp.config",
+    "get_settings": "bipixie_mcp.config",
+    "build_server": "bipixie_mcp.server",
+    "main": "bipixie_mcp.server",
+}
+
+
+def __getattr__(name: str) -> object:
+    """
+    Lazy loader for the public API.
+
+    Importing a public name (e.g. ``from bipixie_mcp import Settings``) triggers
+    this function, which imports the relevant sub-module on first access and caches
+    the result in the package namespace so subsequent accesses are O(1).
+    """
+    if name not in _LAZY_MAP:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+    module_name = _LAZY_MAP[name]
+    import importlib
+
+    module = importlib.import_module(module_name)
+    obj = getattr(module, name)
+
+    # Cache in the package namespace to avoid repeated __getattr__ calls.
+    globals()[name] = obj
+    return obj

bipixie_mcp/auth.py

@@ -0,0 +1,446 @@
+"""Token acquisition for the BI Pixie MCP server.
+
+This module is the single auth boundary for every outbound call the server makes
+to the Power BI REST API. It maps the configured ``BIPIXIE_MCP_AUTH_MODE`` onto a
+concrete ``azure.core.credentials.TokenCredential`` and hands back a cached bearer
+token scoped to ``settings.powerbi_scope``
+(``https://analysis.windows.net/powerbi/api/.default``).
+
+Auth modes are selected by how the server is HOSTED, not by row visibility: the BI
+Pixie semantic model ships no row-level security, so every mode sees the same full
+usage data. Each Enterprise team installs the BI Pixie Dashboard under its own
+license key / container and points its own MCP server at its own model, so a caller
+who can run the server is already entitled to that model's usage data. The same
+package serves both multi-tenant cloud and self-hosted enterprise deployments:
+
+* LOCAL (stdio / Claude Code / Codex on a developer machine)
+    - ``device_code`` (default): :class:`azure.identity.DeviceCodeCredential`
+    - ``interactive``: :class:`azure.identity.InteractiveBrowserCredential`
+      Both are *public-client* flows (no client secret) and persist an MSAL token
+      cache at ``settings.token_cache_path`` so the user is not re-prompted on every
+      restart.
+    - ``azure_cli``: :class:`azure.identity.AzureCliCredential` — reuses an existing
+      ``az login`` with no app registration; the convenient local-dev / validation
+      mode.
+
+* HOSTED (streamable-http inside the customer's Azure)
+    - ``service_principal``: :class:`azure.identity.ClientSecretCredential`
+      (mirrors ``control_plane_app/shared/template_app_client.py``)
+    - ``managed_identity``: :class:`azure.identity.DefaultAzureCredential`
+      (zero-secret; the MI is granted Member/Build on the workspace)
+
+``ResponsePiiFilter`` + ``BIPIXIE_MCP_PII_COLUMNS_ALLOWED`` remain the privacy
+control for user-identifying columns, independent of auth mode.
+
+Security notes
+--------------
+* This module performs READ-ONLY data access only -- it never requests a
+  ``ReadWrite`` scope; the resolved ``.default`` scope should map to the app reg's
+  ``Dataset.Read.All``.
+* Secret values and raw bearer tokens are NEVER logged. A redacting
+  :class:`logging.Filter` (:class:`SecretRedactingFilter`) is attached to this
+  module's logger and masks anything that looks like a GUID, a long hex/base64url
+  blob, or a JWT. Per project convention all log output is emitted to stderr;
+  configuration of handlers/levels happens in ``server.py``.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import re
+import threading
+import time
+from pathlib import Path
+from typing import TYPE_CHECKING, Final
+
+from azure.core.credentials import AccessToken, TokenCredential
+from azure.core.exceptions import ClientAuthenticationError
+from azure.identity import (
+    AzureCliCredential,
+    ClientSecretCredential,
+    DefaultAzureCredential,
+    DeviceCodeCredential,
+    InteractiveBrowserCredential,
+)
+
+if TYPE_CHECKING:  # avoid an import cycle / hard dependency at module-load time
+    from .config import Settings
+
+__all__ = [
+    "AuthError",
+    "TokenProvider",
+    "build_credential",
+    "get_powerbi_token",
+    "SecretRedactingFilter",
+]
+
+
+# --------------------------------------------------------------------------- #
+# Logging (stderr only; handlers configured in server.py). Secrets redacted.
+# --------------------------------------------------------------------------- #
+
+
+class SecretRedactingFilter(logging.Filter):
+    """Mask secret-shaped substrings in any record routed through this logger.
+
+    Defense-in-depth: even if a future code path accidentally interpolates a token,
+    secret, or GUID into a log message, this filter scrubs it before the record
+    reaches a handler. We intentionally over-redact (GUIDs, long hex/base64url runs,
+    JWTs) rather than risk leaking an Entra credential or bearer token.
+    """
+
+    # GUID (tenant/client IDs), 40+ char hex (token fragments / client secrets),
+    # 20+ char base64url runs, and three-segment JWTs.
+    _PATTERNS: Final[tuple[re.Pattern[str], ...]] = (
+        re.compile(r"\b[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-"
+                   r"[0-9a-fA-F]{4}-[0-9a-fA-F]{12}\b"),
+        re.compile(r"\beyJ[A-Za-z0-9_\-]+\.[A-Za-z0-9_\-]+\.[A-Za-z0-9_\-]+"),
+        re.compile(r"\b[0-9a-fA-F]{40,}\b"),
+        re.compile(r"\b[A-Za-z0-9_\-]{20,}\b"),
+    )
+    _MASK: Final[str] = "***REDACTED***"
+
+    def _scrub(self, text: str) -> str:
+        for pattern in self._PATTERNS:
+            text = pattern.sub(self._MASK, text)
+        return text
+
+    def filter(self, record: logging.LogRecord) -> bool:  # noqa: A003 - stdlib name
+        try:
+            message = record.getMessage()
+        except Exception:  # pragma: no cover - never let logging break auth
+            return True
+        scrubbed = self._scrub(message)
+        if scrubbed != message:
+            record.msg = scrubbed
+            record.args = ()
+        return True
+
+
+logger = logging.getLogger(__name__)
+logger.addFilter(SecretRedactingFilter())
+
+
+# --------------------------------------------------------------------------- #
+# Errors
+# --------------------------------------------------------------------------- #
+
+
+class AuthError(Exception):
+    """Raised when a bearer token cannot be acquired for the Power BI API.
+
+    Messages are deliberately non-secret and setup-oriented (e.g. a hint that the
+    'Allow service principals to use Power BI APIs' tenant setting is required, or
+    that the workspace role is missing) so they are safe to surface to an operator
+    or an MCP client.
+    """
+
+
+# How long before the real expiry we proactively refresh, so a token never expires
+# mid-flight on a slow Power BI query. Power BI tokens live ~1h; 5 min is generous.
+_REFRESH_SKEW_SECONDS: Final[int] = 300
+
+# Per-credential-construction hint appended to AuthError messages for 403/consent
+# style failures. Kept terse and secret-free.
+_SP_DISABLED_HINT: Final[str] = (
+    "If using service_principal/managed_identity auth, confirm the Power BI tenant "
+    "setting 'Allow service principals to use Power BI APIs' is enabled and that the "
+    "identity has at least the Member role on the target workspace."
+)
+
+
+# --------------------------------------------------------------------------- #
+# Credential construction
+# --------------------------------------------------------------------------- #
+
+
+def _load_token_cache_seed(token_cache_path: str) -> str | None:
+    """Return the serialized MSAL cache blob for a public-client credential.
+
+    azure-identity's public-client credentials accept a ``cache_persistence_options``
+    arg in newer SDKs, but the broadly-portable contract here is to seed the
+    in-memory MSAL cache from a previously serialized blob and let the caller re-serialize
+    on shutdown. We read the blob if the file exists; absence is not an error
+    (first run). Read failures are downgraded to a warning -- a missing cache only
+    costs the user one extra interactive prompt, never correctness.
+    """
+    path = Path(os.path.expanduser(token_cache_path))
+    if not path.is_file():
+        return None
+    try:
+        return path.read_text(encoding="utf-8")
+    except OSError as exc:
+        # Logs only the cache file PATH + exception class name, never token bytes;
+        # SecretRedactingFilter scrubs defensively. Semgrep heuristic false positive:
+        logger.warning(  # nosemgrep: python.lang.security.audit.logging.logger-credential-leak.python-logger-credential-disclosure
+            "Could not read MSAL token cache at %s (%s); an interactive prompt "
+            "may be required.",
+            path,
+            exc.__class__.__name__,
+        )
+        return None
+
+
+def _build_public_client_credential(
+    settings: "Settings",
+) -> TokenCredential:
+    """Construct a DeviceCode/InteractiveBrowser credential for LOCAL stdio use.
+
+    Both are public-client flows -- no client secret. A persistent MSAL token cache
+    at ``settings.token_cache_path`` is enabled where the installed azure-identity
+    version supports ``TokenCachePersistenceOptions`` so the refresh token survives
+    process restarts; otherwise we fall back to an in-memory cache (one prompt per
+    process) without failing.
+    """
+    if not settings.client_id:
+        raise AuthError(
+            f"auth_mode='{settings.auth_mode}' requires BIPIXIE_MCP_CLIENT_ID "
+            "(a public-client app registration in the configured tenant)."
+        )
+
+    common_kwargs: dict[str, object] = {
+        "tenant_id": settings.tenant_id,
+        "client_id": settings.client_id,
+    }
+
+    # Best-effort persistent cache. Import lazily because the class name/availability
+    # varies across azure-identity releases; a failure here must not break auth.
+    cache_options = None
+    try:  # pragma: no cover - environment-dependent
+        from azure.identity import TokenCachePersistenceOptions
+
+        cache_dir = Path(os.path.expanduser(settings.token_cache_path)).parent
+        cache_dir.mkdir(parents=True, exist_ok=True)
+        cache_options = TokenCachePersistenceOptions(
+            name="bipixie-mcp",
+            allow_unencrypted_storage=True,
+        )
+    except Exception as exc:  # noqa: BLE001 - persistence is optional
+        # Logs only the exception class name, never secrets. Semgrep false positive:
+        logger.debug(  # nosemgrep: python.lang.security.audit.logging.logger-credential-leak.python-logger-credential-disclosure
+            "Persistent token cache unavailable (%s); using in-memory cache.",
+            exc.__class__.__name__,
+        )
+
+    if cache_options is not None:
+        common_kwargs["cache_persistence_options"] = cache_options
+
+    try:
+        if settings.auth_mode == "interactive":
+            return InteractiveBrowserCredential(**common_kwargs)  # type: ignore[arg-type]
+        # default LOCAL mode
+        return DeviceCodeCredential(**common_kwargs)  # type: ignore[arg-type]
+    except TypeError:
+        # Older azure-identity that does not accept cache_persistence_options.
+        common_kwargs.pop("cache_persistence_options", None)
+        if settings.auth_mode == "interactive":
+            return InteractiveBrowserCredential(**common_kwargs)  # type: ignore[arg-type]
+        return DeviceCodeCredential(**common_kwargs)  # type: ignore[arg-type]
+
+
+def build_credential(settings: "Settings") -> TokenCredential:
+    """Map ``settings.auth_mode`` onto a concrete azure-identity credential.
+
+    This is the single place that knows how each mode is wired:
+
+    * ``device_code``      -> :class:`DeviceCodeCredential` (public client)
+    * ``interactive``      -> :class:`InteractiveBrowserCredential` (public client)
+    * ``azure_cli``        -> :class:`AzureCliCredential` (reuses ``az login``; local dev)
+    * ``service_principal``-> :class:`ClientSecretCredential` (app token)
+    * ``managed_identity`` -> :class:`DefaultAzureCredential` (app token)
+
+    Construction is side-effect-light: no token is acquired here. :class:`TokenProvider`
+    drives the actual ``get_token`` call lazily. Raises :class:`AuthError` for an
+    unknown mode or missing required configuration (mirroring, but not replacing,
+    ``Settings.validate_runtime``).
+    """
+    mode = settings.auth_mode
+
+    if mode in ("device_code", "interactive"):
+        return _build_public_client_credential(settings)
+
+    if mode == "service_principal":
+        if not (settings.client_id and settings.client_secret):
+            raise AuthError(
+                "auth_mode='service_principal' requires both "
+                "BIPIXIE_MCP_CLIENT_ID and BIPIXIE_MCP_CLIENT_SECRET."
+            )
+        return ClientSecretCredential(
+            tenant_id=settings.tenant_id,
+            client_id=settings.client_id,
+            client_secret=settings.client_secret,
+        )
+
+    if mode == "azure_cli":
+        # Reuse the developer's existing `az login` — no app registration, no
+        # secret. AzureCliCredential shells out to `az account get-access-token`.
+        # Ideal for local validation/dev. Pass tenant_id when set (newer
+        # azure-identity); fall back gracefully on older versions that lack it.
+        try:
+            return (
+                AzureCliCredential(tenant_id=settings.tenant_id)
+                if settings.tenant_id
+                else AzureCliCredential()
+            )
+        except TypeError:
+            return AzureCliCredential()
+
+    if mode == "managed_identity":
+        # DefaultAzureCredential resolves the Azure-hosted managed identity (and
+        # honors AZURE_CLIENT_ID for a user-assigned MI). No secret needed.
+        return DefaultAzureCredential()
+
+    raise AuthError(
+        f"Unknown BIPIXIE_MCP_AUTH_MODE '{mode}'. Expected one of: "
+        "device_code, interactive, azure_cli, service_principal, managed_identity."
+    )
+
+
+# --------------------------------------------------------------------------- #
+# Token provider (cached + auto-refresh)
+# --------------------------------------------------------------------------- #
+
+
+class TokenProvider:
+    """Acquire and cache a Power BI bearer token for the configured identity.
+
+    A single :class:`TokenProvider` is created once at server startup (in
+    ``server.build_server``) and shared by :class:`powerbi_client.PowerBIClient`.
+    The underlying azure-identity credential is constructed lazily on first use and
+    then reused process-wide; tokens are cached and proactively refreshed about
+    five minutes before expiry, so the credential is *not* invoked per request.
+
+    Thread-safety: ``get_token`` is guarded by a lock because the streamable-http
+    transport may dispatch tool calls concurrently. The hot path (valid cached
+    token) takes the lock only briefly to read the cached value.
+    """
+
+    def __init__(self, settings: "Settings") -> None:
+        self._settings = settings
+        self._scope: str = settings.powerbi_scope
+        self._lock = threading.Lock()
+        self._credential: TokenCredential | None = None
+        self._cached_token: AccessToken | None = None
+
+    # -- public API ------------------------------------------------------- #
+
+    def get_token(self) -> str:
+        """Return a valid bearer token for ``settings.powerbi_scope``.
+
+        Reuses the cached token until it is within ``_REFRESH_SKEW_SECONDS`` of
+        expiry, then transparently refreshes. Raises :class:`AuthError` (never a
+        raw azure exception, and never a message containing a secret) on failure.
+        """
+        with self._lock:
+            cached = self._cached_token
+            if cached is not None and not self._is_expiring(cached):
+                return cached.token
+
+            credential = self._ensure_credential()
+            token = self._acquire(credential)
+            self._cached_token = token
+            return token.token
+
+    # -- internals -------------------------------------------------------- #
+
+    @staticmethod
+    def _is_expiring(token: AccessToken) -> bool:
+        """True if the token is unset, already expired, or inside the refresh skew."""
+        return token.expires_on - _REFRESH_SKEW_SECONDS <= int(time.time())
+
+    def _ensure_credential(self) -> TokenCredential:
+        if self._credential is None:
+            self._credential = build_credential(self._settings)
+            # Logs only the auth-mode enum (device_code/azure_cli/...), not a secret.
+            logger.info(  # nosemgrep: python.lang.security.audit.logging.logger-credential-leak.python-logger-credential-disclosure
+                "Initialized Power BI credential (auth_mode=%s).",
+                self._settings.auth_mode,
+            )
+        return self._credential
+
+    def _acquire(self, credential: TokenCredential) -> AccessToken:
+        try:
+            token = credential.get_token(self._scope)
+        except ClientAuthenticationError as exc:
+            raise AuthError(self._explain(exc)) from exc
+        except Exception as exc:  # noqa: BLE001 - normalize to AuthError
+            raise AuthError(
+                f"Failed to acquire a Power BI token ({exc.__class__.__name__}). "
+                f"{_SP_DISABLED_HINT}"
+            ) from exc
+
+        if not token or not getattr(token, "token", None):
+            raise AuthError(
+                "Token acquisition returned an empty bearer for scope "
+                f"'{self._scope}'. {_SP_DISABLED_HINT}"
+            )
+        # Logs only the token EXPIRY timestamp, never the token itself. False positive:
+        logger.debug("Acquired Power BI token (expires_on=%s).", token.expires_on)  # nosemgrep: python.lang.security.audit.logging.logger-credential-leak.python-logger-credential-disclosure
+        return token
+
+    def _explain(self, exc: ClientAuthenticationError) -> str:
+        """Build a non-secret, setup-oriented message from an auth failure.
+
+        We surface the credential type and a remediation hint rather than the raw
+        exception text (which can embed correlation IDs / claims). Service-principal
+        and managed-identity failures get the tenant-setting/workspace-role hint.
+        """
+        mode = self._settings.auth_mode
+        base = (
+            f"Could not authenticate to Power BI using auth_mode='{mode}' "
+            f"for scope '{self._scope}'."
+        )
+        if mode in ("service_principal", "managed_identity"):
+            return f"{base} {_SP_DISABLED_HINT}"
+        if mode in ("device_code", "interactive"):
+            return (
+                f"{base} Confirm BIPIXIE_MCP_CLIENT_ID is a public-client app "
+                "registration in the configured tenant and that you completed the "
+                "sign-in prompt."
+            )
+        if mode == "azure_cli":
+            return (
+                f"{base} Confirm you have run 'az login' and that the active Azure "
+                "CLI account can reach the Power BI API in this tenant."
+            )
+        return base
+
+
+# --------------------------------------------------------------------------- #
+# Module-level async convenience
+# --------------------------------------------------------------------------- #
+
+# Process-wide TokenProvider cache keyed by the identity-relevant config tuple, so
+# repeated get_powerbi_token() calls reuse one credential + token cache. Keyed on the
+# fields that determine the identity/scope so a re-configured Settings yields a fresh
+# provider rather than a stale token.
+_PROVIDER_LOCK = threading.Lock()
+_PROVIDERS: dict[tuple, TokenProvider] = {}
+
+
+def _provider_key(settings: "Settings") -> tuple:
+    return (
+        settings.tenant_id,
+        settings.auth_mode,
+        settings.client_id,
+        settings.powerbi_scope,
+    )
+
+
+async def get_powerbi_token(settings: "Settings") -> str:
+    """Async-friendly accessor returning a cached Power BI bearer token.
+
+    Thin coroutine wrapper over :class:`TokenProvider` for callers (and tools) that
+    operate in an async context. The MSAL/azure-identity token acquisition itself is
+    synchronous and almost always served from cache; the rare network refresh is
+    brief, so it runs inline. Providers are cached per identity-config so concurrent
+    callers share one credential and one token.
+    """
+    key = _provider_key(settings)
+    with _PROVIDER_LOCK:
+        provider = _PROVIDERS.get(key)
+        if provider is None:
+            provider = TokenProvider(settings)
+            _PROVIDERS[key] = provider
+    return provider.get_token()