Authorization Feature

libs/agent_framework/pyproject.toml

@@ -23,7 +23,8 @@ dependencies = [
   "aiohttp>=3.9.0",
   "motor>=3.6.0",
   "google-cloud-pubsub>=2.28.0",
-  "mcp>=1.9.0"
+  "mcp>=1.9.0",
+  "PyJWT[crypto]>=2.9.0"
 ]
 
 [tool.setuptools.packages.find]

libs/agent_framework/src/agent_framework/security/__init__.py

@@ -0,0 +1,40 @@
+from .authentication import (
+    ApiKeyAuthenticationProvider,
+    AuthenticatedPrincipal,
+    AuthenticationProvider,
+    AuthenticationResult,
+    BasicAuthenticationProvider,
+    DenyAuthenticationProvider,
+    JwtAuthenticationProvider,
+    NoAuthenticationProvider,
+    OAuth2IntrospectionAuthenticationProvider,
+    StaticBearerAuthenticationProvider,
+    TrustedProxyAuthenticationProvider,
+    verify_secret,
+)
+from .factory import create_authentication_provider, create_provider_from_config, env_provider_config
+from .installer import install_authentication, load_authentication_policies
+from .middleware import AuthenticationMiddleware, AuthenticationPolicy, PolicyAuthenticationMiddleware
+
+__all__ = [
+    "ApiKeyAuthenticationProvider",
+    "AuthenticatedPrincipal",
+    "AuthenticationProvider",
+    "AuthenticationResult",
+    "AuthenticationMiddleware",
+    "AuthenticationPolicy",
+    "BasicAuthenticationProvider",
+    "DenyAuthenticationProvider",
+    "JwtAuthenticationProvider",
+    "NoAuthenticationProvider",
+    "OAuth2IntrospectionAuthenticationProvider",
+    "PolicyAuthenticationMiddleware",
+    "StaticBearerAuthenticationProvider",
+    "TrustedProxyAuthenticationProvider",
+    "create_authentication_provider",
+    "create_provider_from_config",
+    "env_provider_config",
+    "install_authentication",
+    "load_authentication_policies",
+    "verify_secret",
+]

libs/agent_framework/src/agent_framework/security/authentication.py

@@ -0,0 +1,190 @@
+from __future__ import annotations
+
+import base64
+import hashlib
+import hmac
+import logging
+import time
+from dataclasses import dataclass, field
+from typing import Any, Mapping, Protocol, Sequence
+
+import httpx
+from fastapi import Request
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class AuthenticatedPrincipal:
+    subject: str
+    scheme: str
+    claims: Mapping[str, Any] = field(default_factory=dict)
+
+
+@dataclass(frozen=True)
+class AuthenticationResult:
+    authenticated: bool
+    principal: AuthenticatedPrincipal | None = None
+    error: str | None = None
+    challenge: str | None = None
+
+
+class AuthenticationProvider(Protocol):
+    async def authenticate(self, request: Request) -> AuthenticationResult: ...
+
+
+def _constant_time_equals(left: str, right: str) -> bool:
+    return hmac.compare_digest(left.encode("utf-8"), right.encode("utf-8"))
+
+
+def _pbkdf2_hash(secret: str, salt: str, iterations: int = 310_000) -> str:
+    digest = hashlib.pbkdf2_hmac("sha256", secret.encode(), salt.encode(), iterations)
+    return base64.urlsafe_b64encode(digest).decode().rstrip("=")
+
+
+def verify_secret(secret: str, stored_value: str) -> bool:
+    """Accepts plain:<value>, sha256:<hex>, or pbkdf2_sha256:<iterations>:<salt>:<digest>."""
+    if stored_value.startswith("plain:"):
+        return _constant_time_equals(secret, stored_value.removeprefix("plain:"))
+    if stored_value.startswith("sha256:"):
+        candidate = hashlib.sha256(secret.encode()).hexdigest()
+        return _constant_time_equals(candidate, stored_value.removeprefix("sha256:"))
+    if stored_value.startswith("pbkdf2_sha256:"):
+        try:
+            _, iterations, salt, expected = stored_value.split(":", 3)
+            return _constant_time_equals(_pbkdf2_hash(secret, salt, int(iterations)), expected)
+        except (ValueError, TypeError):
+            return False
+    return _constant_time_equals(secret, stored_value)
+
+
+class NoAuthenticationProvider:
+    async def authenticate(self, request: Request) -> AuthenticationResult:
+        return AuthenticationResult(True, AuthenticatedPrincipal("anonymous", "none"))
+
+
+class DenyAuthenticationProvider:
+    async def authenticate(self, request: Request) -> AuthenticationResult:
+        return AuthenticationResult(False, error="authentication_policy_not_configured")
+
+
+class BasicAuthenticationProvider:
+    def __init__(self, client_id: str, secret_hash: str, realm: str = "agent-api"):
+        self.client_id = client_id
+        self.secret_hash = secret_hash
+        self.realm = realm
+
+    async def authenticate(self, request: Request) -> AuthenticationResult:
+        header = request.headers.get("authorization", "")
+        if not header.lower().startswith("basic "):
+            return AuthenticationResult(False, error="missing_basic_credentials", challenge=f'Basic realm="{self.realm}"')
+        try:
+            decoded = base64.b64decode(header.split(" ", 1)[1], validate=True).decode("utf-8")
+            supplied_id, supplied_secret = decoded.split(":", 1)
+        except (ValueError, UnicodeDecodeError):
+            return AuthenticationResult(False, error="invalid_basic_credentials", challenge=f'Basic realm="{self.realm}"')
+        valid = _constant_time_equals(supplied_id, self.client_id) and verify_secret(supplied_secret, self.secret_hash)
+        if not valid:
+            return AuthenticationResult(False, error="invalid_basic_credentials", challenge=f'Basic realm="{self.realm}"')
+        return AuthenticationResult(True, AuthenticatedPrincipal(supplied_id, "basic"))
+
+
+class ApiKeyAuthenticationProvider:
+    def __init__(self, expected_hash: str, header_name: str = "x-api-key", principal: str = "api-client"):
+        self.expected_hash = expected_hash
+        self.header_name = header_name.lower()
+        self.principal = principal
+
+    async def authenticate(self, request: Request) -> AuthenticationResult:
+        supplied = request.headers.get(self.header_name)
+        if not supplied or not verify_secret(supplied, self.expected_hash):
+            return AuthenticationResult(False, error="invalid_api_key")
+        return AuthenticationResult(True, AuthenticatedPrincipal(self.principal, "api_key"))
+
+
+class StaticBearerAuthenticationProvider:
+    def __init__(self, token_hash: str, principal: str = "bearer-client"):
+        self.token_hash = token_hash
+        self.principal = principal
+
+    async def authenticate(self, request: Request) -> AuthenticationResult:
+        header = request.headers.get("authorization", "")
+        if not header.lower().startswith("bearer "):
+            return AuthenticationResult(False, error="missing_bearer_token", challenge="Bearer")
+        token = header.split(" ", 1)[1]
+        if not verify_secret(token, self.token_hash):
+            return AuthenticationResult(False, error="invalid_bearer_token", challenge="Bearer")
+        return AuthenticationResult(True, AuthenticatedPrincipal(self.principal, "bearer"))
+
+
+class JwtAuthenticationProvider:
+    def __init__(self, key: str, algorithms: Sequence[str], audience: str | None = None, issuer: str | None = None):
+        try:
+            import jwt  # type: ignore
+        except ImportError as exc:
+            raise RuntimeError("JWT authentication requires PyJWT[crypto]") from exc
+        self.jwt = jwt
+        self.key = key
+        self.algorithms = list(algorithms)
+        self.audience = audience
+        self.issuer = issuer
+
+    async def authenticate(self, request: Request) -> AuthenticationResult:
+        header = request.headers.get("authorization", "")
+        if not header.lower().startswith("bearer "):
+            return AuthenticationResult(False, error="missing_bearer_token", challenge="Bearer")
+        token = header.split(" ", 1)[1]
+        try:
+            claims = self.jwt.decode(token, self.key, algorithms=self.algorithms, audience=self.audience, issuer=self.issuer)
+        except Exception as exc:
+            logger.info("JWT rejected: %s", exc.__class__.__name__)
+            return AuthenticationResult(False, error="invalid_jwt", challenge="Bearer")
+        subject = str(claims.get("sub") or claims.get("client_id") or "jwt-client")
+        return AuthenticationResult(True, AuthenticatedPrincipal(subject, "jwt", claims))
+
+
+class OAuth2IntrospectionAuthenticationProvider:
+    def __init__(self, introspection_url: str, client_id: str, client_secret: str, timeout_seconds: float = 5.0):
+        self.introspection_url = introspection_url
+        self.client_id = client_id
+        self.client_secret = client_secret
+        self.timeout_seconds = timeout_seconds
+
+    async def authenticate(self, request: Request) -> AuthenticationResult:
+        header = request.headers.get("authorization", "")
+        if not header.lower().startswith("bearer "):
+            return AuthenticationResult(False, error="missing_bearer_token", challenge="Bearer")
+        token = header.split(" ", 1)[1]
+        try:
+            async with httpx.AsyncClient(timeout=self.timeout_seconds) as client:
+                response = await client.post(
+                    self.introspection_url,
+                    data={"token": token},
+                    auth=(self.client_id, self.client_secret),
+                    headers={"accept": "application/json"},
+                )
+                response.raise_for_status()
+                claims = response.json()
+        except (httpx.HTTPError, ValueError):
+            return AuthenticationResult(False, error="introspection_unavailable", challenge="Bearer")
+        if not claims.get("active") or (claims.get("exp") and int(claims["exp"]) <= int(time.time())):
+            return AuthenticationResult(False, error="inactive_token", challenge="Bearer")
+        subject = str(claims.get("sub") or claims.get("client_id") or claims.get("username") or "oauth-client")
+        return AuthenticationResult(True, AuthenticatedPrincipal(subject, "oauth2_introspection", claims))
+
+
+class TrustedProxyAuthenticationProvider:
+    def __init__(self, subject_header: str = "x-authenticated-subject", shared_secret_header: str | None = None, shared_secret_hash: str | None = None):
+        self.subject_header = subject_header.lower()
+        self.shared_secret_header = shared_secret_header.lower() if shared_secret_header else None
+        self.shared_secret_hash = shared_secret_hash
+
+    async def authenticate(self, request: Request) -> AuthenticationResult:
+        subject = request.headers.get(self.subject_header)
+        if not subject:
+            return AuthenticationResult(False, error="missing_trusted_subject")
+        if self.shared_secret_header and self.shared_secret_hash:
+            supplied = request.headers.get(self.shared_secret_header)
+            if not supplied or not verify_secret(supplied, self.shared_secret_hash):
+                return AuthenticationResult(False, error="invalid_proxy_signature")
+        return AuthenticationResult(True, AuthenticatedPrincipal(subject, "trusted_proxy"))

libs/agent_framework/src/agent_framework/security/factory.py

@@ -0,0 +1,112 @@
+from __future__ import annotations
+
+import os
+from collections.abc import Mapping
+from typing import Any
+
+from .authentication import (
+    ApiKeyAuthenticationProvider,
+    BasicAuthenticationProvider,
+    DenyAuthenticationProvider,
+    JwtAuthenticationProvider,
+    NoAuthenticationProvider,
+    OAuth2IntrospectionAuthenticationProvider,
+    StaticBearerAuthenticationProvider,
+    TrustedProxyAuthenticationProvider,
+)
+
+
+def _required_env(name: str) -> str:
+    value = os.getenv(name)
+    if value is None or not value.strip():
+        raise ValueError(f"Required authentication environment variable is missing: {name}")
+    return value
+
+
+def _resolve(config: Mapping[str, Any], key: str, *, required: bool = False, default: Any = None) -> Any:
+    env_key = config.get(f"{key}_env")
+    if env_key:
+        value = os.getenv(str(env_key))
+        if required and (value is None or not value.strip()):
+            raise ValueError(f"Required authentication environment variable is missing: {env_key}")
+        return value if value is not None else default
+    value = config.get(key, default)
+    if required and (value is None or (isinstance(value, str) and not value.strip())):
+        raise ValueError(f"Required authentication configuration is missing: {key}")
+    return value
+
+
+def create_provider_from_config(config: Mapping[str, Any]):
+    """Create a provider from a secret-safe mapping.
+
+    Secret values may be supplied indirectly with ``<field>_env`` keys so YAML
+    never needs to contain credentials.
+    """
+    mode = str(config.get("mode", "none")).strip().lower()
+    if mode in {"none", "disabled"}:
+        return NoAuthenticationProvider()
+    if mode in {"deny", "reject"}:
+        return DenyAuthenticationProvider()
+    if mode == "basic":
+        return BasicAuthenticationProvider(
+            str(_resolve(config, "client_id", required=True)),
+            str(_resolve(config, "secret_hash", required=True)),
+            str(_resolve(config, "realm", default="agent-api")),
+        )
+    if mode == "api_key":
+        return ApiKeyAuthenticationProvider(
+            str(_resolve(config, "api_key_hash", required=True)),
+            str(_resolve(config, "header", default="x-api-key")),
+            str(_resolve(config, "principal", default="api-client")),
+        )
+    if mode == "bearer_static":
+        return StaticBearerAuthenticationProvider(
+            str(_resolve(config, "token_hash", required=True)),
+            str(_resolve(config, "principal", default="bearer-client")),
+        )
+    if mode == "jwt":
+        algorithms = _resolve(config, "algorithms", default=["RS256"])
+        if isinstance(algorithms, str):
+            algorithms = [item.strip() for item in algorithms.split(",") if item.strip()]
+        return JwtAuthenticationProvider(
+            str(_resolve(config, "key", required=True)),
+            algorithms,
+            _resolve(config, "audience"),
+            _resolve(config, "issuer"),
+        )
+    if mode == "oauth2_introspection":
+        return OAuth2IntrospectionAuthenticationProvider(
+            str(_resolve(config, "introspection_url", required=True)),
+            str(_resolve(config, "client_id", required=True)),
+            str(_resolve(config, "client_secret", required=True)),
+            float(_resolve(config, "timeout_seconds", default=5)),
+        )
+    if mode == "trusted_proxy":
+        return TrustedProxyAuthenticationProvider(
+            str(_resolve(config, "subject_header", default="x-authenticated-subject")),
+            _resolve(config, "shared_secret_header"),
+            _resolve(config, "shared_secret_hash"),
+        )
+    raise ValueError(f"Unsupported authentication mode: {mode}")
+
+
+def env_provider_config(prefix: str = "AGENT_AUTH") -> dict[str, Any]:
+    mode = os.getenv(f"{prefix}_MODE", "none").strip().lower()
+    config: dict[str, Any] = {"mode": mode}
+    if mode == "basic":
+        config.update(client_id=_required_env(f"{prefix}_BASIC_CLIENT_ID"), secret_hash=_required_env(f"{prefix}_BASIC_SECRET_HASH"), realm=os.getenv(f"{prefix}_BASIC_REALM", "agent-api"))
+    elif mode == "api_key":
+        config.update(api_key_hash=_required_env(f"{prefix}_API_KEY_HASH"), header=os.getenv(f"{prefix}_API_KEY_HEADER", "x-api-key"), principal=os.getenv(f"{prefix}_API_KEY_PRINCIPAL", "api-client"))
+    elif mode == "bearer_static":
+        config.update(token_hash=_required_env(f"{prefix}_BEARER_TOKEN_HASH"), principal=os.getenv(f"{prefix}_BEARER_PRINCIPAL", "bearer-client"))
+    elif mode == "jwt":
+        config.update(key=_required_env(f"{prefix}_JWT_KEY"), algorithms=os.getenv(f"{prefix}_JWT_ALGORITHMS", "RS256"), audience=os.getenv(f"{prefix}_JWT_AUDIENCE") or None, issuer=os.getenv(f"{prefix}_JWT_ISSUER") or None)
+    elif mode == "oauth2_introspection":
+        config.update(introspection_url=_required_env(f"{prefix}_OAUTH2_INTROSPECTION_URL"), client_id=_required_env(f"{prefix}_OAUTH2_CLIENT_ID"), client_secret=_required_env(f"{prefix}_OAUTH2_CLIENT_SECRET"), timeout_seconds=float(os.getenv(f"{prefix}_OAUTH2_TIMEOUT_SECONDS", "5")))
+    elif mode == "trusted_proxy":
+        config.update(subject_header=os.getenv(f"{prefix}_PROXY_SUBJECT_HEADER", "x-authenticated-subject"), shared_secret_header=os.getenv(f"{prefix}_PROXY_SHARED_SECRET_HEADER") or None, shared_secret_hash=os.getenv(f"{prefix}_PROXY_SHARED_SECRET_HASH") or None)
+    return config
+
+
+def create_authentication_provider(prefix: str = "AGENT_AUTH"):
+    return create_provider_from_config(env_provider_config(prefix))

libs/agent_framework/src/agent_framework/security/installer.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Any
+
+import yaml
+from fastapi import FastAPI
+
+from .authentication import DenyAuthenticationProvider
+from .factory import create_authentication_provider, create_provider_from_config
+from .middleware import AuthenticationMiddleware, AuthenticationPolicy, PolicyAuthenticationMiddleware
+
+
+def _csv(value: str | None, default: str = "") -> list[str]:
+    return [item.strip() for item in (value if value is not None else default).split(",") if item.strip()]
+
+
+def _bool(value: str | None, default: bool = False) -> bool:
+    if value is None:
+        return default
+    return value.strip().lower() in {"1", "true", "yes", "on"}
+
+
+def load_authentication_policies(path: str | Path) -> tuple[list[AuthenticationPolicy], Any]:
+    raw = yaml.safe_load(Path(path).read_text(encoding="utf-8")) or {}
+    providers = {
+        name: create_provider_from_config(config or {})
+        for name, config in (raw.get("providers") or {}).items()
+    }
+    policies: list[AuthenticationPolicy] = []
+    for index, item in enumerate(raw.get("policies") or []):
+        provider_name = item.get("provider")
+        if provider_name not in providers:
+            raise ValueError(f"Unknown authentication provider in policy: {provider_name}")
+        policies.append(AuthenticationPolicy(
+            name=str(item.get("name") or f"policy-{index + 1}"),
+            provider=providers[provider_name],
+            paths=tuple(item.get("paths") or ["*"]),
+            methods=frozenset(str(method).upper() for method in (item.get("methods") or [])),
+            required_roles=frozenset(str(role) for role in (item.get("required_roles") or [])),
+            required_scopes=frozenset(str(scope) for scope in (item.get("required_scopes") or [])),
+        ))
+    default_name = raw.get("default_provider")
+    default_provider = providers.get(default_name) if default_name else DenyAuthenticationProvider()
+    return policies, default_provider
+
+
+def install_authentication(app: FastAPI, prefix: str = "AGENT_AUTH") -> bool:
+    """Install optional authentication using an isolated environment prefix.
+
+    Returns True when middleware was installed. Authentication remains disabled
+    unless ``<PREFIX>_ENABLED=true`` or a non-``none`` mode/policy file is set.
+    """
+    policy_file = os.getenv(f"{prefix}_POLICIES_FILE")
+    mode = os.getenv(f"{prefix}_MODE", "none").strip().lower()
+    enabled = _bool(os.getenv(f"{prefix}_ENABLED"), default=bool(policy_file or mode not in {"none", "disabled"}))
+    if not enabled:
+        return False
+
+    if policy_file:
+        policies, default_provider = load_authentication_policies(policy_file)
+        app.add_middleware(PolicyAuthenticationMiddleware, policies=policies, default_provider=default_provider)
+        return True
+
+    provider = create_authentication_provider(prefix)
+    public_paths = _csv(os.getenv(f"{prefix}_PUBLIC_PATHS"), "/health,/ready,/live,/docs,/openapi.json,/redoc")
+    public_prefixes = _csv(os.getenv(f"{prefix}_PUBLIC_PREFIXES"))
+    app.add_middleware(AuthenticationMiddleware, provider=provider, public_paths=public_paths, public_prefixes=public_prefixes)
+    return True

libs/agent_framework/src/agent_framework/security/middleware.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+import fnmatch
+import logging
+from collections.abc import Iterable, Sequence
+from dataclasses import dataclass, field
+
+from fastapi import Request
+from fastapi.responses import JSONResponse
+from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
+from starlette.responses import Response
+
+from .authentication import AuthenticationProvider, DenyAuthenticationProvider
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class AuthenticationPolicy:
+    name: str
+    provider: AuthenticationProvider
+    paths: tuple[str, ...] = ("*",)
+    methods: frozenset[str] = field(default_factory=frozenset)
+    required_roles: frozenset[str] = field(default_factory=frozenset)
+    required_scopes: frozenset[str] = field(default_factory=frozenset)
+
+    def matches(self, path: str, method: str) -> bool:
+        method_matches = not self.methods or method.upper() in self.methods
+        return method_matches and any(fnmatch.fnmatchcase(path, pattern) for pattern in self.paths)
+
+
+def _claim_values(claims, names: Sequence[str]) -> set[str]:
+    values: set[str] = set()
+    for name in names:
+        raw = claims.get(name)
+        if isinstance(raw, str):
+            values.update(item for item in raw.replace(",", " ").split() if item)
+        elif isinstance(raw, (list, tuple, set)):
+            values.update(str(item) for item in raw)
+    return values
+
+
+class AuthenticationMiddleware(BaseHTTPMiddleware):
+    """Backward-compatible single-provider middleware."""
+
+    def __init__(self, app, provider: AuthenticationProvider, public_paths: Iterable[str] = (), public_prefixes: Iterable[str] = ()):
+        super().__init__(app)
+        self.provider = provider
+        self.public_paths = frozenset(public_paths)
+        self.public_prefixes = tuple(public_prefixes)
+
+    def _is_public(self, path: str) -> bool:
+        return path in self.public_paths or any(path.startswith(prefix) for prefix in self.public_prefixes)
+
+    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
+        if request.method == "OPTIONS" or self._is_public(request.url.path):
+            return await call_next(request)
+        return await _authenticate_request(request, call_next, self.provider)
+
+
+class PolicyAuthenticationMiddleware(BaseHTTPMiddleware):
+    """Selects the first matching route policy and authenticates the request."""
+
+    def __init__(self, app, policies: Sequence[AuthenticationPolicy], default_provider: AuthenticationProvider | None = None):
+        super().__init__(app)
+        self.policies = tuple(policies)
+        self.default_provider = default_provider or DenyAuthenticationProvider()
+
+    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
+        if request.method == "OPTIONS":
+            return await call_next(request)
+        policy = next((item for item in self.policies if item.matches(request.url.path, request.method)), None)
+        if policy is None:
+            return await _authenticate_request(request, call_next, self.default_provider)
+        return await _authenticate_request(
+            request,
+            call_next,
+            policy.provider,
+            policy_name=policy.name,
+            required_roles=policy.required_roles,
+            required_scopes=policy.required_scopes,
+        )
+
+
+async def _authenticate_request(request: Request, call_next: RequestResponseEndpoint, provider: AuthenticationProvider, *, policy_name: str | None = None, required_roles: frozenset[str] = frozenset(), required_scopes: frozenset[str] = frozenset()) -> Response:
+    result = await provider.authenticate(request)
+    if not result.authenticated or result.principal is None:
+        headers = {"WWW-Authenticate": result.challenge} if result.challenge else None
+        return JSONResponse(status_code=401, content={"detail": "Unauthorized", "code": result.error or "unauthorized", "policy": policy_name}, headers=headers)
+
+    roles = _claim_values(result.principal.claims, ("roles", "role", "groups"))
+    scopes = _claim_values(result.principal.claims, ("scope", "scp", "scopes"))
+    if required_roles and not required_roles.issubset(roles):
+        return JSONResponse(status_code=403, content={"detail": "Forbidden", "code": "missing_required_role", "policy": policy_name})
+    if required_scopes and not required_scopes.issubset(scopes):
+        return JSONResponse(status_code=403, content={"detail": "Forbidden", "code": "missing_required_scope", "policy": policy_name})
+
+    request.state.auth_principal = result.principal
+    request.state.auth_policy = policy_name
+    return await call_next(request)