Source code for fsh_lib.oauth.providers

from __future__ import annotations

from typing import TYPE_CHECKING, Literal

import httpx
from pydantic import BaseModel, ConfigDict

if TYPE_CHECKING:
    from collections.abc import Sequence

TokenAuthMethod = Literal["client_secret_post", "client_secret_basic"]

_DEFAULT_AUTH_METHOD: TokenAuthMethod = "client_secret_post"
_DEFAULT_TIMEOUT = 10.0

GOOGLE_ISSUER = "https://accounts.google.com"
_MS_WILDCARD_TENANTS = frozenset({"common", "organizations", "consumers"})




class OAuthConfigError(ValueError):
    pass





class ClientCredentials(BaseModel):
    model_config = ConfigDict(frozen=True)

    client_id: str
    client_secret_env: str | None = None
    redirect_path: str = ""





class ProviderConfig(BaseModel):
    model_config = ConfigDict(frozen=True, extra="ignore")

    authorization_endpoint: str
    token_endpoint: str
    userinfo_endpoint: str | None = None
    jwks_uri: str | None = None
    issuer: str | None = None
    token_endpoint_auth_method: TokenAuthMethod = _DEFAULT_AUTH_METHOD
    default_scopes: tuple[str, ...] = ()
    verify_email: bool = False

    tenant: str | None = None
    tenant_base: str | None = None

    def for_tenant(self, tenant: str) -> ProviderConfig:
        if self.tenant is None or self.tenant_base is None:
            msg = "provider endpoints are not tenant-scoped"
            raise OAuthConfigError(msg)

        base = f"{self.tenant_base.rstrip('/')}/{tenant}"
        issuer = None if tenant in _MS_WILDCARD_TENANTS else f"{base}/v2.0"

        return self.model_copy(
            update={
                "authorization_endpoint": f"{base}/oauth2/v2.0/authorize",
                "token_endpoint": f"{base}/oauth2/v2.0/token",
                "jwks_uri": f"{base}/discovery/v2.0/keys",
                "issuer": issuer,
                "tenant": tenant,
            },
        )

    @property
    def admin_consent_endpoint(self) -> str:
        if self.tenant is None or self.tenant_base is None:
            msg = "provider endpoints are not tenant-scoped"
            raise OAuthConfigError(msg)

        base = f"{self.tenant_base.rstrip('/')}/{self.tenant}"

        return f"{base}/adminconsent"



async def discover(
    issuer: str,
    *,
    client: httpx.AsyncClient | None = None,
    default_scopes: Sequence[str] = ("openid", "email", "profile"),
    verify_email: bool = False,
    timeout: float = _DEFAULT_TIMEOUT,  # noqa: ASYNC109 -- forwarded to httpx
) -> ProviderConfig:
    url = f"{issuer.rstrip('/')}/.well-known/openid-configuration"
    owns_client = client is None
    client = client or httpx.AsyncClient(timeout=timeout)

    try:
        response = await client.get(url)
        response.raise_for_status()
        document = response.json()

    except httpx.HTTPError as exc:
        msg = f"OIDC discovery failed for {issuer!r}: {exc}"
        raise OAuthConfigError(msg) from exc

    except ValueError as exc:
        msg = f"OIDC discovery document for {issuer!r} is not valid JSON"
        raise OAuthConfigError(msg) from exc

    finally:
        if owns_client:
            await client.aclose()

    missing = [
        key
        for key in ("authorization_endpoint", "token_endpoint")
        if not document.get(key)
    ]

    if missing:
        msg = (
            f"OIDC discovery document for {issuer!r} is missing "
            f"required field(s): {missing}"
        )
        raise OAuthConfigError(msg)

    document_issuer = document.get("issuer")

    if document_issuer is not None and document_issuer.rstrip(
        "/"
    ) != issuer.rstrip("/"):
        msg = (
            f"OIDC discovery document issuer {document_issuer!r} does "
            f"not match the requested issuer {issuer!r}"
        )
        raise OAuthConfigError(msg)

    return ProviderConfig(
        authorization_endpoint=document["authorization_endpoint"],
        token_endpoint=document["token_endpoint"],
        userinfo_endpoint=document.get("userinfo_endpoint"),
        jwks_uri=document.get("jwks_uri"),
        issuer=document_issuer or issuer,
        default_scopes=tuple(default_scopes),
        verify_email=verify_email,
    )


def google(
    *,
    default_scopes: Sequence[str] = ("openid", "email", "profile"),
    verify_email: bool = False,
) -> ProviderConfig:
    return ProviderConfig(
        authorization_endpoint="https://accounts.google.com/o/oauth2/v2/auth",
        token_endpoint="https://oauth2.googleapis.com/token",  # noqa: S106
        userinfo_endpoint="https://openidconnect.googleapis.com/v1/userinfo",
        jwks_uri="https://www.googleapis.com/oauth2/v3/certs",
        issuer=GOOGLE_ISSUER,
        default_scopes=tuple(default_scopes),
        verify_email=verify_email,
    )


def microsoft(
    tenant: str = "common",
    *,
    default_scopes: Sequence[str] = ("openid", "email", "profile"),
    verify_email: bool = False,
) -> ProviderConfig:
    base = f"https://login.microsoftonline.com/{tenant}"
    issuer = None if tenant in _MS_WILDCARD_TENANTS else f"{base}/v2.0"

    return ProviderConfig(
        authorization_endpoint=f"{base}/oauth2/v2.0/authorize",
        token_endpoint=f"{base}/oauth2/v2.0/token",
        userinfo_endpoint="https://graph.microsoft.com/oidc/userinfo",
        jwks_uri=f"{base}/discovery/v2.0/keys",
        issuer=issuer,
        default_scopes=tuple(default_scopes),
        verify_email=verify_email,
        tenant=tenant,
        tenant_base="https://login.microsoftonline.com",
    )