pesco-ncr/backend/app/auth/entra.py

"""Entra ID access-token validation (python-jose + tenant JWKS)."""
import logging
import time

import httpx
from jose import JWTError, jwt

from app.config import get_settings

logger = logging.getLogger(__name__)


class AuthError(Exception):
    def __init__(self, message: str, status_code: int = 401):
        self.message = message
        self.status_code = status_code
        super().__init__(message)


_jwks: dict[str, dict] = {}
_jwks_fetched_at: float = 0.0
_JWKS_TTL = 60 * 60 * 12


async def _fetch_jwks() -> None:
    global _jwks, _jwks_fetched_at
    settings = get_settings()
    url = (
        f"https://login.microsoftonline.com/{settings.entra_tenant_id}"
        "/discovery/v2.0/keys"
    )
    async with httpx.AsyncClient(timeout=15) as client:
        resp = await client.get(url)
        resp.raise_for_status()
        _jwks = {k["kid"]: k for k in resp.json().get("keys", [])}
        _jwks_fetched_at = time.time()
    logger.info("Fetched %d Entra signing keys", len(_jwks))


async def _get_signing_key(kid: str) -> dict:
    stale = time.time() - _jwks_fetched_at > _JWKS_TTL
    if kid not in _jwks or stale:
        await _fetch_jwks()
    key = _jwks.get(kid)
    if key is None:
        raise AuthError("Token signed with an unknown key.")
    return key


async def validate_access_token(token: str) -> dict:
    """Validate signature, expiry, audience, and issuer; return claims."""
    settings = get_settings()
    if not settings.entra_tenant_id or not settings.entra_client_id:
        raise AuthError(
            "Entra ID is not configured (ENTRA_TENANT_ID / ENTRA_CLIENT_ID).", 503
        )
    try:
        header = jwt.get_unverified_header(token)
        key = await _get_signing_key(header.get("kid", ""))
        claims = jwt.decode(
            token,
            key,
            algorithms=["RS256"],
            options={"verify_aud": False},  # audience list checked below
        )
    except AuthError:
        raise
    except JWTError as exc:
        raise AuthError(f"Invalid token: {exc}") from exc

    aud = claims.get("aud")
    if aud not in settings.api_audiences:
        raise AuthError("Token audience does not match this API.")

    tid = settings.entra_tenant_id
    valid_issuers = {
        f"https://login.microsoftonline.com/{tid}/v2.0",
        f"https://sts.windows.net/{tid}/",
    }
    if claims.get("iss") not in valid_issuers:
        raise AuthError("Token issuer does not match the configured tenant.")
    return claims


async def ensure_group_membership(claims: dict, token: str) -> None:
    """Front-door gate: require membership in ENTRA_ALLOWED_GROUP_ID.

    Uses the `groups` claim when present; on claim overage falls back to a
    delegated Graph checkMemberGroups call.
    """
    settings = get_settings()
    group_id = settings.entra_allowed_group_id
    if not group_id:
        return  # gate disabled by configuration

    groups = claims.get("groups")
    if groups is not None:
        if group_id in groups:
            return
        raise AuthError(
            "Your account is not a member of the NCR access group.", 403
        )

    claim_names = claims.get("_claim_names") or {}
    if "groups" in claim_names:
        # Group overage: too many groups to embed in the token.
        from app.services.graph import check_member_group

        try:
            if await check_member_group(token, group_id):
                return
        except Exception as exc:
            logger.warning("Group overage Graph check failed: %s", exc)
            raise AuthError(
                "Could not verify group membership (group overage). Ensure the "
                "app has delegated GroupMember.Read.All consent, or scope the "
                "group claim to 'Groups assigned to the application'.", 403
            ) from exc
        raise AuthError("Your account is not a member of the NCR access group.", 403)

    raise AuthError(
        "Access token has no groups claim. Add the groups claim in the app "
        "registration (Token configuration → Add groups claim).", 403
    )