Source code for ska_aaa_authhelpers.token_scheme

from collections.abc import AsyncGenerator, Iterable, Mapping
from typing import Annotated

from fastapi import Depends, Header
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from fastapi.security.base import SecurityBase
from pydantic import JsonValue
from starlette_context import request_cycle_context

from . import validation_rules as vr
from .auth_context import AuthContext
from .errors import AuthFailError
from .jwt import KeysType, decode_jwt, issued_by_microsoft
from .roles import ENTRA_ROLE_GRANTING_GROUPS, INDIGO_ROLE_GRANTING_GROUPS, Role

USER_ID_CLAIM = "preferred_username"

access_token = HTTPBearer(
    scheme_name="SKAO Authentication",
    description=(
        "OAuth2 access token authorization via Indigo IAM "
        "(current) or Microsoft EntraID (deprecated)."
    ),
    bearerFormat="JWT Access Token",
    auto_error=True,
)


# TODO: Remove once fully migrated off Entra ID.
def entra_id_auth_context(
    claims: dict,
    access_token: str,
    trace: str,
) -> AuthContext:
    groups = frozenset(claims.get("groups", []))
    tkn_roles = frozenset(claims.get("roles", [])).intersection(Role)
    grp_roles = frozenset(
        ENTRA_ROLE_GRANTING_GROUPS[g] for g in groups if g in ENTRA_ROLE_GRANTING_GROUPS
    )
    roles = tkn_roles | grp_roles | frozenset({Role.ANY})
    if scp := claims.get("scp"):  # User tokens have scopes
        user_id = claims["oid"]
        scopes = frozenset(scp.split(" "))
    else:  # Client credentials tokens do not.
        scopes = frozenset()
        user_id = claims["appid"] if claims.get("ver") == "1.0" else claims["azp"]
    return AuthContext(
        audience=claims["aud"],
        user_id=user_id,
        trace=trace,
        groups=groups,
        scopes=scopes,
        principals=groups.union((user_id,)),
        roles=roles,
        token_claims=claims,
        access_token=access_token,
    )


def is_client_credentials(claims: Mapping[str, JsonValue]):
    # If the 'sub'ject of the token is the client_id, then this is a
    # client credentials token issued directly to a client, not
    # to a human user:
    return claims["sub"] == claims.get("client_id")


def get_user_id_and_inferred_roles(claims: Mapping[str, JsonValue]) -> tuple[str, frozenset[Role]]:
    if is_client_credentials(claims):
        return str(claims["client_id"]), frozenset({Role.ANY, Role.APP2APP})
    elif user_id := claims.get(USER_ID_CLAIM):
        return str(user_id), frozenset({Role.ANY})
    else:
        msg = f"Access token lacks '{USER_ID_CLAIM}' -- clients must request the 'profile' scope."
        raise AuthFailError(msg)


def get_auth_context(
    authorization: HTTPAuthorizationCredentials,
    issuers: Iterable[str],
    audiences: Iterable[str],
    keys: KeysType,
    trace: str | None = None,
) -> AuthContext:
    trace = trace if trace is not None else ""
    claims = decode_jwt(authorization.credentials, keys=keys, issuers=issuers, audiences=audiences)

    if issued_by_microsoft(claims["iss"]):
        return entra_id_auth_context(claims, authorization.credentials, trace)

    groups = frozenset(claims.get("groups", []))
    grp_roles = frozenset(
        INDIGO_ROLE_GRANTING_GROUPS[g] for g in groups if g in INDIGO_ROLE_GRANTING_GROUPS
    )
    user_id, nongroup_roles = get_user_id_and_inferred_roles(claims)
    roles = grp_roles | nongroup_roles
    scopes = frozenset(claims.get("scope", "").split())

    return AuthContext(
        audience=claims["aud"],
        user_id=user_id,
        trace=trace,
        groups=groups,
        scopes=scopes,
        principals=groups.union((user_id,)),
        roles=roles,
        token_claims=claims,
        access_token=authorization.credentials,
    )


[docs] class TokenScheme(SecurityBase): """ Security scheme that generates an AuthContext from bearer token and and enforces basic authorisation rules. """ def __init__( self, *, audiences: frozenset[str], roles: frozenset[Role], scopes: frozenset[str], groups: frozenset[str], app_ids: frozenset[str], issuers: Iterable[str], keys: KeysType, ): self.scopes = scopes self.roles = roles self.groups = groups self.app_ids = app_ids self.audiences = audiences self.issuers = issuers self.scheme_name = access_token.scheme_name self.model = access_token.model self.keys = keys async def __call__( self, authorization: Annotated[HTTPAuthorizationCredentials, Depends(access_token)], X_Request_ID: Annotated[str | None, Header()] = None, ) -> AsyncGenerator[AuthContext, None]: # https://starlette-context.readthedocs.io/en/latest/fastapi.html # Create a global request-scoped context object that # we rely on to populate auth context into log messages. # See audit_log_filter.py. ac = get_auth_context( authorization, self.issuers, self.audiences, self.keys, trace=X_Request_ID ) with request_cycle_context({"auth": ac}): vr.roles_contain(self.roles, ac.roles) if Role.APP2APP in ac.roles: # For app2pp calls, we check for an exact match of allowed app IDs. vr.app_id_matches(ac.user_id, self.app_ids) else: # For user calls, we check the user has granted this necessary scopes. vr.contains(self.scopes, ac.scopes, label="scopes") if self.groups: vr.contains(self.groups, ac.groups, label="groups") yield ac