orchestrata_ultimate.gateway.auth

Reference › orchestrata_ultimate › gateway › auth

This module holds the three authenticators behind the gateway's GATEWAY_AUTH_MODE setting (SPEC-014 FR-005 to FR-007). auth_handler chooses one per gateway, auth handlers use it for the auth.* RPCs, and the kernel CLI client and the gateway token-minting script also import it. Settings come from config; failures raise AuthenticationError from exceptions. See also gateway.

• TokenAuthenticator is the jwt mode. create_token issues a signed token for a user id and optional claims, with its lifetime set by the TOKEN_EXPIRY_SECONDS setting; validate_token decodes it and rejects expired or invalid tokens and tokens with no sub claim.
• StaticTokenAuthenticator is the static mode. GATEWAY_TOKEN only seeds the token; the live value is read from Redis, and rotate replaces it so every gateway sharing that Redis picks it up on the next attempt without a restart. authenticate rejects a blank token, compares in constant time, and returns claims with sub and a read_write scope.
• HmacAuthenticator is the hmac mode. compute_signature builds an HMAC-SHA256 over timestamp, nonce and client id. authenticate checks the timestamp against the GATEWAY_HMAC_DRIFT_SECONDS window, accepts a signature matching any configured secret (so old and new secrets can overlap during rotation), and records the nonce in Redis with an atomic set-if-absent so a replay is refused. The scope comes from GATEWAY_HMAC_CLIENT_SCOPES, with a default scope for unlisted clients.

The Redis clients are created lazily through core.redis_factory.

At a glance

Class diagram

Dependency diagram

API

Classes

TokenAuthenticator

Line 19

JWT token authenticator for gateway connections.

StaticTokenAuthenticator

Decorators: traces(spec='014', story='US-006', reqs=['REQ-014-091', 'REQ-014-092', 'REQ-014-093', 'REQ-014-094', 'REQ-014-095']) · Line 120

Pre-shared static token authenticator (SPEC-014 FR-006).

HmacAuthenticator

Decorators: traces(spec='014', story='US-006', reqs=['REQ-014-096', 'REQ-014-097', 'REQ-014-098', 'REQ-014-099', 'REQ-014-100']) · Line 236

HMAC-SHA256 request authenticator (SPEC-014 FR-007).

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/gateway/auth.py

"""Token-based authentication for WebSocket connections."""

import hashlib
import hmac
import time
from typing import Any

import jwt
import redis.asyncio as redis

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.core.redis_factory import create_async_redis
from orchestrata_ultimate.gateway.config import GatewayConfig
from orchestrata_ultimate.gateway.exceptions import (
    AuthenticationError,
)


class TokenAuthenticator:
    """JWT token authenticator for gateway connections."""

    def __init__(self, config: GatewayConfig) -> None:
        """Initialise from gateway configuration."""
        self._secret = config.jwt_secret
        self._algorithm = config.jwt_algorithm
        self._expiry = config.token_expiry_seconds

    @staticmethod
    def _get_secret() -> str:
        """Return the JWT secret the gateway signs and validates with.

        Resolves it through :class:`GatewayConfig` (the ``JWT_SECRET``
        environment variable, read at call time, falling back to the same
        ``dev-secret`` default the gateway uses). Out-of-band token issuers —
        e.g. the demo token minted in ``scripts/start_test_services.py`` — call
        this so they sign with exactly the secret the running gateway will
        validate against, instead of duplicating the env-lookup and drifting.
        """
        return GatewayConfig().jwt_secret

    def create_token(
        self,
        user_id: str,
        claims: dict[str, Any] | None = None,
    ) -> str:
        """Create a signed JWT for *user_id*.

        Args:
            user_id: Subject identifier for the token.
            claims: Optional extra claims merged into
                the JWT payload.

        Returns:
            Encoded JWT string.

        Raises:
            TypeError: If user_id is not a str or claims
                is not a dict.
            ValueError: If user_id is empty.
        """
        if not isinstance(user_id, str):
            raise TypeError("user_id must be a str")
        if not user_id:
            raise ValueError("user_id must not be empty")
        if claims is not None and not isinstance(claims, dict):
            raise TypeError("claims must be a dict or None")
        now = int(time.time())
        payload: dict[str, Any] = {
            "sub": user_id,
            "iat": now,
            "exp": now + self._expiry,
        }
        if claims:
            payload.update(claims)
        return jwt.encode(
            payload,
            self._secret,
            algorithm=self._algorithm,
        )

    def validate_token(self, token: str) -> dict[str, Any]:
        """Validate and decode a JWT token.

        Args:
            token: Encoded JWT string.

        Returns:
            Decoded payload dictionary.

        Raises:
            AuthenticationError: On expired or invalid
                token, or missing 'sub' claim.
        """
        try:
            payload = jwt.decode(
                token,
                self._secret,
                algorithms=[self._algorithm],
            )
        except jwt.ExpiredSignatureError as exc:
            raise AuthenticationError("Token has expired") from exc
        except jwt.InvalidTokenError as exc:
            raise AuthenticationError(f"Invalid token: {exc}") from exc
        if "sub" not in payload:
            raise AuthenticationError("Token missing 'sub' claim")
        return payload


@traces(
    spec="014",
    story="US-006",
    reqs=[
        "REQ-014-091",
        "REQ-014-092",
        "REQ-014-093",
        "REQ-014-094",
        "REQ-014-095",
    ],
)
class StaticTokenAuthenticator:
    """Pre-shared static token authenticator (SPEC-014 FR-006).

    The token is configured via :class:`GatewayConfig` / the ``GATEWAY_TOKEN``
    environment variable (REQ-014-091), but that config/env value is only the
    SEED: the live token is Redis-backed (REQ-014-094) so a deployment can
    call :meth:`rotate` to change it and every gateway process sharing that
    Redis instance picks it up on the next auth attempt -- no restart. A
    blank/missing client token is rejected before any comparison
    (REQ-014-093); a configured token is compared with
    ``hmac.compare_digest`` for constant-time equality (REQ-014-092).
    """

    #: Scope granted on successful static-token auth (REQ-014-095).
    DEFAULT_SCOPE = "read_write"

    #: Redis key holding the live (rotatable) token. Callers that need
    #: isolation -- tests, primarily -- may override via *redis_key*.
    _DEFAULT_REDIS_KEY = "gw:auth:static_token"

    def __init__(self, config: GatewayConfig, redis_key: str | None = None) -> None:
        """Initialise from gateway configuration.

        Args:
            config: Gateway configuration; supplies the Redis connection
                parameters and the config/env-seeded token.
            redis_key: Override for the Redis key the live token is stored
                under. Defaults to a fixed, process-wide key; tests pass a
                unique key to avoid cross-test interference in a shared
                Redis database.
        """
        self._config = config
        self._seed_token = config.static_token
        self._redis_key = redis_key or self._DEFAULT_REDIS_KEY
        self._client: redis.Redis | None = None

    async def _get_client(self) -> redis.Redis:
        """Lazily construct (and cache) the async Redis client."""
        if self._client is None:
            self._client = create_async_redis(
                host=self._config.redis_host,
                port=self._config.redis_port,
                db=self._config.redis_db,
                decode_responses=True,
            )
        return self._client

    async def current_token(self) -> str | None:
        """Return the live token.

        The Redis-stored value wins when present (it reflects the most
        recent :meth:`rotate` call); otherwise falls back to the
        config/env-seeded value that constructed this authenticator.
        """
        client = await self._get_client()
        stored = await client.get(self._redis_key)
        if stored:
            # redis-py types get() as Any. _get_client builds the client with
            # decode_responses=True, so a hit is already `str`; decoding
            # defensively keeps a future flip of that flag from handing back
            # the repr of a bytes object ("b'...'") as if it were the token.
            return stored.decode() if isinstance(stored, bytes) else str(stored)
        return self._seed_token

    async def rotate(self, new_token: str) -> None:
        """Rotate the live token to *new_token* (REQ-014-094).

        Writes the new value to Redis; every gateway process that shares
        this Redis instance validates against it starting with the very
        next auth attempt -- no gateway restart required.

        Raises:
            ValueError: *new_token* is blank.
        """
        if not new_token:
            raise ValueError("new_token must not be empty")
        client = await self._get_client()
        await client.set(self._redis_key, new_token)

    async def authenticate(self, token: str) -> dict[str, Any]:
        """Validate *token* against the live static token.

        Args:
            token: Client-provided token.

        Returns:
            Claims dict with ``sub`` and ``scope`` (mirrors
            :meth:`TokenAuthenticator.validate_token`'s shape so callers can
            treat the two authenticators uniformly).

        Raises:
            AuthenticationError: *token* is blank/missing (REQ-014-093), no
                static token is configured (neither Redis nor config/env),
                or *token* does not match the live token.
        """
        if not token:
            raise AuthenticationError("Missing or blank token")
        current = await self.current_token()
        if not current:
            raise AuthenticationError("Static token auth is not configured")
        if not hmac.compare_digest(token, current):
            raise AuthenticationError("Invalid token")
        return {"sub": "static-token", "scope": self.DEFAULT_SCOPE}


@traces(
    spec="014",
    story="US-006",
    reqs=[
        "REQ-014-096",
        "REQ-014-097",
        "REQ-014-098",
        "REQ-014-099",
        "REQ-014-100",
    ],
)
class HmacAuthenticator:
    """HMAC-SHA256 request authenticator (SPEC-014 FR-007).

    Validates ``HMAC-SHA256(shared_secret, timestamp + nonce + client_id)``
    (REQ-014-096) against a configurable timestamp-drift window (default
    60s -- REQ-014-097), rejects a replayed nonce within that window via a
    Redis-tracked TTL (REQ-014-098), and tries every configured shared
    secret in turn so multiple secrets can be live at once for key rotation
    (REQ-014-099). Scope comes from :attr:`GatewayConfig.hmac_client_scopes`
    -- a client_id -> scope mapping -- falling back to :attr:`DEFAULT_SCOPE`
    for an unmapped client_id (REQ-014-100).
    """

    #: Fallback scope for a client_id with no entry in hmac_client_scopes.
    DEFAULT_SCOPE = "read_only"

    _NONCE_REDIS_KEY_PREFIX = "gw:auth:hmac:nonce"

    def __init__(self, config: GatewayConfig) -> None:
        """Initialise from gateway configuration."""
        self._config = config
        self._secrets = config.hmac_secrets
        self._drift_seconds = config.hmac_drift_seconds
        self._scopes = config.hmac_client_scopes
        self._client: redis.Redis | None = None

    async def _get_client(self) -> redis.Redis:
        """Lazily construct (and cache) the async Redis client."""
        if self._client is None:
            self._client = create_async_redis(
                host=self._config.redis_host,
                port=self._config.redis_port,
                db=self._config.redis_db,
                decode_responses=True,
            )
        return self._client

    @staticmethod
    def compute_signature(
        secret: str,
        timestamp: str,
        nonce: str,
        client_id: str,
    ) -> str:
        """Compute ``HMAC-SHA256(secret, timestamp + nonce + client_id)``.

        Returns:
            Hex-encoded digest (REQ-014-096).
        """
        message = f"{timestamp}{nonce}{client_id}".encode()
        return hmac.new(secret.encode(), message, hashlib.sha256).hexdigest()

    async def authenticate(
        self,
        client_id: str,
        timestamp: str,
        nonce: str,
        signature: str,
    ) -> dict[str, Any]:
        """Validate an HMAC-signed authentication request.

        Args:
            client_id: Claimed client identifier.
            timestamp: Unix timestamp (seconds, as a string) the request
                was signed at.
            nonce: Per-request random value; replay-checked.
            signature: Hex-encoded ``HMAC-SHA256(secret, timestamp + nonce
                + client_id)`` the client computed.

        Returns:
            Claims dict with ``sub`` (the client_id) and ``scope``.

        Raises:
            AuthenticationError: A required field is missing, no shared
                secrets are configured, the timestamp drift exceeds the
                configured window (REQ-014-097), the signature does not
                match any configured secret (REQ-014-096/099), or the nonce
                has already been used within the drift window
                (REQ-014-098).
        """
        if not client_id or not timestamp or not nonce or not signature:
            raise AuthenticationError(
                "Missing client_id, timestamp, nonce, or signature"
            )
        if not self._secrets:
            raise AuthenticationError("HMAC auth is not configured (no shared secrets)")
        try:
            request_time = float(timestamp)
        except ValueError as exc:
            raise AuthenticationError("Invalid timestamp") from exc
        drift = abs(time.time() - request_time)
        if drift > self._drift_seconds:
            raise AuthenticationError(
                f"Timestamp drift {drift:.1f}s exceeds window of "
                f"{self._drift_seconds}s"
            )
        # First configured secret whose signature matches wins -- supports
        # rotation where an old and new secret are both briefly live.
        if not any(
            hmac.compare_digest(
                signature,
                self.compute_signature(secret, timestamp, nonce, client_id),
            )
            for secret in self._secrets
        ):
            raise AuthenticationError("Invalid HMAC signature")
        client = await self._get_client()
        nonce_key = f"{self._NONCE_REDIS_KEY_PREFIX}:{client_id}:{nonce}"
        # Atomic check-and-set: the nonce is only ever "seen" once, and it
        # expires with the drift window it is scoped to (a timestamp older
        # than that window is already rejected above, so nothing outlives
        # its relevance).
        first_use = await client.set(
            nonce_key,
            "1",
            nx=True,
            ex=max(1, self._drift_seconds),
        )
        if not first_use:
            raise AuthenticationError("Replayed nonce")
        scope = self._scopes.get(client_id, self.DEFAULT_SCOPE)
        return {"sub": client_id, "scope": scope}