byoi 0.1.0a1__py3-none-any.whl

byoi/telemetry.py

@@ -0,0 +1,352 @@
+"""OpenTelemetry tracing support for BYOI.
+
+This module provides optional OpenTelemetry instrumentation for BYOI operations.
+Tracing can help with debugging, monitoring, and understanding the performance
+of authentication flows.
+
+Installation:
+    pip install byoi[telemetry]
+
+Usage:
+    from byoi.telemetry import configure_tracing, get_tracer
+
+    # Configure tracing (call once at startup)
+    configure_tracing(service_name="my-auth-service")
+
+    # Or use with custom TracerProvider
+    from opentelemetry.sdk.trace import TracerProvider
+    provider = TracerProvider()
+    configure_tracing(tracer_provider=provider)
+"""
+
+from contextlib import contextmanager
+from functools import wraps
+from typing import TYPE_CHECKING, Any, Callable, TypeVar
+
+# Type variable for generic function decoration
+F = TypeVar("F", bound=Callable[..., Any])
+
+# Flag to track if OpenTelemetry is available
+_OTEL_AVAILABLE = False
+_tracer = None
+
+try:
+    from opentelemetry import trace
+    from opentelemetry.trace import Span, SpanKind, Status, StatusCode, Tracer
+    from opentelemetry.sdk.trace import TracerProvider
+    from opentelemetry.sdk.resources import Resource
+    from opentelemetry.semconv.resource import ResourceAttributes
+
+    _OTEL_AVAILABLE = True
+except ImportError:
+    # OpenTelemetry not installed - provide stub implementations
+    trace = None  # type: ignore[assignment]
+    Span = None  # type: ignore[assignment, misc]
+    SpanKind = None  # type: ignore[assignment, misc]
+    Status = None  # type: ignore[assignment, misc]
+    StatusCode = None  # type: ignore[assignment, misc]
+    Tracer = None  # type: ignore[assignment, misc]
+    TracerProvider = None  # type: ignore[assignment, misc]
+    Resource = None  # type: ignore[assignment, misc]
+    ResourceAttributes = None  # type: ignore[assignment, misc]
+
+
+__all__ = (
+    "is_tracing_available",
+    "configure_tracing",
+    "get_tracer",
+    "traced",
+    "traced_async",
+    "span_context",
+)
+
+
+# BYOI-specific span names
+SPAN_CREATE_AUTH_URL = "byoi.create_authorization_url"
+SPAN_EXCHANGE_CODE = "byoi.exchange_code"
+SPAN_REFRESH_TOKEN = "byoi.refresh_token"
+SPAN_AUTHENTICATE = "byoi.authenticate"
+SPAN_LINK_IDENTITY = "byoi.link_identity"
+SPAN_UNLINK_IDENTITY = "byoi.unlink_identity"
+SPAN_VALIDATE_TOKEN = "byoi.validate_token"
+SPAN_DISCOVER_PROVIDER = "byoi.discover_provider"
+SPAN_FETCH_JWKS = "byoi.fetch_jwks"
+
+
+def is_tracing_available() -> bool:
+    """Check if OpenTelemetry tracing is available.
+
+    Returns:
+        True if OpenTelemetry is installed and configured, False otherwise.
+    """
+    return _OTEL_AVAILABLE
+
+
+def configure_tracing(
+    service_name: str = "byoi",
+    tracer_provider: "TracerProvider | None" = None,
+    **resource_attributes: str,
+) -> "Tracer | None":
+    """Configure OpenTelemetry tracing for BYOI.
+
+    This function should be called once during application startup to
+    enable tracing for BYOI operations.
+
+    Args:
+        service_name: The name of the service for tracing.
+        tracer_provider: Optional custom TracerProvider. If not provided,
+            a new one will be created.
+        **resource_attributes: Additional resource attributes to include.
+
+    Returns:
+        The configured Tracer, or None if OpenTelemetry is not available.
+
+    Example:
+        ```python
+        from byoi.telemetry import configure_tracing
+
+        # Simple configuration
+        configure_tracing(service_name="my-auth-service")
+
+        # With custom attributes
+        configure_tracing(
+            service_name="my-auth-service",
+            environment="production",
+            version="1.0.0",
+        )
+        ```
+    """
+    global _tracer
+
+    if not _OTEL_AVAILABLE:
+        return None
+
+    if tracer_provider is None:
+        # Create a new TracerProvider with resource attributes
+        attributes = {
+            ResourceAttributes.SERVICE_NAME: service_name,
+            **resource_attributes,
+        }
+        resource = Resource.create(attributes)
+        tracer_provider = TracerProvider(resource=resource)
+        trace.set_tracer_provider(tracer_provider)
+
+    _tracer = trace.get_tracer(
+        instrumenting_module_name="byoi",
+        instrumenting_library_version="0.1.0a1",
+    )
+
+    return _tracer
+
+
+def get_tracer() -> "Tracer | None":
+    """Get the configured BYOI tracer.
+
+    Returns:
+        The configured Tracer, or None if tracing is not configured.
+    """
+    global _tracer
+
+    if not _OTEL_AVAILABLE:
+        return None
+
+    if _tracer is None:
+        # Try to get from global TracerProvider
+        _tracer = trace.get_tracer(
+            instrumenting_module_name="byoi",
+            instrumenting_library_version="0.1.0a1",
+        )
+
+    return _tracer
+
+
+@contextmanager
+def span_context(
+    name: str,
+    kind: "SpanKind | None" = None,
+    attributes: dict[str, Any] | None = None,
+):
+    """Context manager for creating spans.
+
+    This is a convenience wrapper that handles the case where tracing
+    is not available.
+
+    Args:
+        name: The name of the span.
+        kind: The kind of span (CLIENT, SERVER, INTERNAL, etc.).
+        attributes: Optional attributes to add to the span.
+
+    Yields:
+        The created Span, or a no-op context if tracing is not available.
+
+    Example:
+        ```python
+        from byoi.telemetry import span_context
+
+        with span_context("my_operation", attributes={"user_id": "123"}) as span:
+            # Do work
+            span.set_attribute("result", "success")
+        ```
+    """
+    if not _OTEL_AVAILABLE:
+        yield _NoOpSpan()
+        return
+
+    tracer = get_tracer()
+    if tracer is None:
+        yield _NoOpSpan()
+        return
+
+    span_kind = kind or SpanKind.INTERNAL
+    with tracer.start_as_current_span(
+        name,
+        kind=span_kind,
+        attributes=attributes,
+    ) as span:
+        yield span
+
+
+class _NoOpSpan:
+    """A no-op span for when tracing is not available."""
+
+    def set_attribute(self, key: str, value: Any) -> None:
+        """No-op set_attribute."""
+        pass
+
+    def set_attributes(self, attributes: dict[str, Any]) -> None:
+        """No-op set_attributes."""
+        pass
+
+    def add_event(self, name: str, attributes: dict[str, Any] | None = None) -> None:
+        """No-op add_event."""
+        pass
+
+    def record_exception(self, exception: BaseException) -> None:
+        """No-op record_exception."""
+        pass
+
+    def set_status(self, status: Any) -> None:
+        """No-op set_status."""
+        pass
+
+
+def traced(
+    span_name: str | None = None,
+    kind: "SpanKind | None" = None,
+    record_exception: bool = True,
+) -> Callable[[F], F]:
+    """Decorator to trace a synchronous function.
+
+    Args:
+        span_name: The name of the span. Defaults to the function name.
+        kind: The kind of span.
+        record_exception: Whether to record exceptions in the span.
+
+    Returns:
+        A decorator function.
+
+    Example:
+        ```python
+        from byoi.telemetry import traced
+
+        @traced("my_operation")
+        def my_function():
+            # Do work
+            pass
+        ```
+    """
+
+    def decorator(func: F) -> F:
+        if not _OTEL_AVAILABLE:
+            return func
+
+        @wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> Any:
+            name = span_name or func.__name__
+            with span_context(name, kind=kind) as span:
+                try:
+                    result = func(*args, **kwargs)
+                    return result
+                except Exception as e:
+                    if record_exception:
+                        span.record_exception(e)
+                        span.set_status(Status(StatusCode.ERROR, str(e)))
+                    raise
+
+        return wrapper  # type: ignore[return-value]
+
+    return decorator
+
+
+def traced_async(
+    span_name: str | None = None,
+    kind: "SpanKind | None" = None,
+    record_exception: bool = True,
+) -> Callable[[F], F]:
+    """Decorator to trace an asynchronous function.
+
+    Args:
+        span_name: The name of the span. Defaults to the function name.
+        kind: The kind of span.
+        record_exception: Whether to record exceptions in the span.
+
+    Returns:
+        A decorator function.
+
+    Example:
+        ```python
+        from byoi.telemetry import traced_async
+
+        @traced_async("my_async_operation")
+        async def my_async_function():
+            # Do async work
+            pass
+        ```
+    """
+
+    def decorator(func: F) -> F:
+        if not _OTEL_AVAILABLE:
+            return func
+
+        @wraps(func)
+        async def wrapper(*args: Any, **kwargs: Any) -> Any:
+            name = span_name or func.__name__
+            with span_context(name, kind=kind) as span:
+                try:
+                    result = await func(*args, **kwargs)
+                    return result
+                except Exception as e:
+                    if record_exception:
+                        span.record_exception(e)
+                        span.set_status(Status(StatusCode.ERROR, str(e)))
+                    raise
+
+        return wrapper  # type: ignore[return-value]
+
+    return decorator
+
+
+def add_auth_attributes(
+    span: "Span | _NoOpSpan",
+    provider_name: str | None = None,
+    user_id: str | None = None,
+    identity_id: str | None = None,
+    is_new_user: bool | None = None,
+) -> None:
+    """Add common authentication attributes to a span.
+
+    Args:
+        span: The span to add attributes to.
+        provider_name: The name of the OIDC provider.
+        user_id: The user ID.
+        identity_id: The linked identity ID.
+        is_new_user: Whether this is a new user.
+    """
+    if provider_name is not None:
+        span.set_attribute("byoi.provider_name", provider_name)
+    if user_id is not None:
+        span.set_attribute("byoi.user_id", user_id)
+    if identity_id is not None:
+        span.set_attribute("byoi.identity_id", identity_id)
+    if is_new_user is not None:
+        span.set_attribute("byoi.is_new_user", is_new_user)

byoi/tokens.py

@@ -0,0 +1,340 @@
+"""ID Token validation utilities.
+
+Handles JWT validation including signature verification, claims validation,
+and nonce checking.
+"""
+
+import base64
+import hashlib
+import logging
+from typing import Any
+
+from jose import JWTError, jwt
+from jose.exceptions import ExpiredSignatureError, JWTClaimsError
+
+from byoi.errors import (
+    InvalidAudienceError,
+    InvalidIssuerError,
+    InvalidNonceError,
+    TokenExpiredError,
+    TokenValidationError,
+)
+from byoi.models import IdentityInfo, OIDCProviderConfig
+from byoi.providers import ProviderManager
+
+logger = logging.getLogger("byoi.tokens")
+
+__all__ = (
+    "TokenValidator",
+    "validate_id_token",
+)
+
+
+class TokenValidator:
+    """Validates ID tokens from OIDC providers."""
+
+    # Standard OIDC claims to extract
+    STANDARD_CLAIMS = {
+        "sub",
+        "email",
+        "email_verified",
+        "name",
+        "given_name",
+        "family_name",
+        "picture",
+        "locale",
+        "iss",
+        "aud",
+        "exp",
+        "iat",
+        "nonce",
+        "at_hash",
+        "c_hash",
+        "auth_time",
+        "acr",
+        "amr",
+        "azp",
+    }
+
+    def __init__(self, provider_manager: ProviderManager) -> None:
+        """Initialize the token validator.
+
+        Args:
+            provider_manager: The provider manager for fetching JWKS.
+        """
+        self._provider_manager = provider_manager
+
+    async def validate(
+        self,
+        id_token: str,
+        provider_name: str,
+        expected_nonce: str | None = None,
+        verify_at_hash: bool = False,
+        access_token: str | None = None,
+    ) -> IdentityInfo:
+        """Validate an ID token and extract identity information.
+
+        Args:
+            id_token: The raw ID token (JWT).
+            provider_name: The provider name.
+            expected_nonce: The expected nonce value (for replay protection).
+            verify_at_hash: Whether to verify the at_hash claim.
+            access_token: The access token (required if verify_at_hash is True).
+
+        Returns:
+            Extracted identity information.
+
+        Raises:
+            TokenValidationError: If validation fails.
+            ProviderNotFoundError: If the provider is not registered.
+        """
+        logger.debug("Validating ID token for provider=%s", provider_name)
+
+        config = self._provider_manager.get_provider(provider_name)
+        jwks = await self._provider_manager.get_jwks(provider_name)
+
+        # Decode and validate the token, with JWKS refresh retry on key not found
+        try:
+            claims = await self._decode_and_validate(id_token, config, jwks)
+        except TokenValidationError as e:
+            # Check if this might be a key rotation issue (kid not found in JWKS)
+            error_msg = str(e).lower()
+            if "key" in error_msg or "kid" in error_msg or "signature" in error_msg:
+                logger.info(
+                    "Token validation failed, retrying with fresh JWKS for provider=%s",
+                    provider_name,
+                )
+                # Fetch fresh JWKS bypassing cache
+                jwks = await self._provider_manager.get_jwks(provider_name, force_refresh=True)
+                # Retry validation with fresh JWKS
+                claims = await self._decode_and_validate(id_token, config, jwks)
+            else:
+                raise
+
+        # Validate nonce if provided
+        if expected_nonce is not None:
+            token_nonce = claims.get("nonce")
+            if token_nonce != expected_nonce:
+                logger.warning(
+                    "Nonce mismatch for provider=%s: expected=%s, got=%s",
+                    provider_name,
+                    expected_nonce,
+                    token_nonce,
+                )
+                raise InvalidNonceError()
+
+        # Validate at_hash if required
+        if verify_at_hash and access_token:
+            await self._verify_at_hash(claims, access_token, id_token)
+
+        # Extract identity information
+        identity = self._extract_identity(claims, provider_name)
+
+        logger.debug(
+            "ID token validated successfully: provider=%s, subject=%s",
+            provider_name,
+            identity.subject,
+        )
+
+        return identity
+
+    async def _decode_and_validate(
+        self,
+        id_token: str,
+        config: OIDCProviderConfig,
+        jwks: dict[str, Any],
+    ) -> dict[str, Any]:
+        """Decode and validate the JWT.
+
+        Args:
+            id_token: The raw ID token.
+            config: The provider configuration.
+            jwks: The provider's JWKS.
+
+        Returns:
+            The validated claims.
+
+        Raises:
+            TokenValidationError: If validation fails.
+        """
+        # Determine expected audience
+        audience = config.audience or config.client_id
+
+        try:
+            # First, decode without verification to get the header
+            unverified = jwt.get_unverified_header(id_token)
+            algorithm = unverified.get("alg", "RS256")
+
+            # Decode and verify the token
+            claims = jwt.decode(
+                id_token,
+                jwks,
+                algorithms=[algorithm],
+                audience=audience,
+                issuer=config.issuer,
+                options={
+                    "verify_aud": True,
+                    "verify_iss": True,
+                    "verify_exp": True,
+                    "verify_iat": True,
+                    "verify_nbf": False,  # Not all providers include nbf
+                    "leeway": config.clock_skew_seconds,
+                },
+            )
+        except ExpiredSignatureError as e:
+            logger.warning("Token expired for provider=%s", config.name)
+            raise TokenExpiredError() from e
+        except JWTClaimsError as e:
+            error_msg = str(e).lower()
+            if "audience" in error_msg:
+                # Try to extract actual audience from token
+                try:
+                    unverified_claims = jwt.get_unverified_claims(id_token)
+                    actual_aud = unverified_claims.get("aud", "unknown")
+                except Exception:
+                    actual_aud = "unknown"
+                logger.warning(
+                    "Audience mismatch for provider=%s: expected=%s, got=%s",
+                    config.name,
+                    audience,
+                    actual_aud,
+                )
+                raise InvalidAudienceError(audience, actual_aud) from e
+            if "issuer" in error_msg:
+                try:
+                    unverified_claims = jwt.get_unverified_claims(id_token)
+                    actual_iss = unverified_claims.get("iss", "unknown")
+                except Exception:
+                    actual_iss = "unknown"
+                logger.warning(
+                    "Issuer mismatch for provider=%s: expected=%s, got=%s",
+                    config.name,
+                    config.issuer,
+                    actual_iss,
+                )
+                raise InvalidIssuerError(config.issuer, actual_iss) from e
+            logger.warning("Token claims validation failed for provider=%s: %s", config.name, str(e))
+            raise TokenValidationError(f"Token claims validation failed: {e}") from e
+        except JWTError as e:
+            logger.error("JWT validation failed for provider=%s: %s", config.name, str(e))
+            raise TokenValidationError(f"Token validation failed: {e}") from e
+
+        # Validate required claims
+        if "sub" not in claims:
+            logger.error("Token missing required 'sub' claim for provider=%s", config.name)
+            raise TokenValidationError("Token missing required 'sub' claim")
+
+        return claims
+
+    async def _verify_at_hash(
+        self,
+        claims: dict[str, Any],
+        access_token: str,
+        id_token: str,
+    ) -> None:
+        """Verify the at_hash claim matches the access token.
+
+        Args:
+            claims: The ID token claims.
+            access_token: The access token.
+            id_token: The ID token (to determine algorithm).
+
+        Raises:
+            TokenValidationError: If at_hash verification fails.
+        """
+
+        at_hash = claims.get("at_hash")
+        if at_hash is None:
+            # at_hash is optional, but if verify_at_hash is True and it's missing,
+            # we should raise an error
+            raise TokenValidationError("Token missing 'at_hash' claim")
+
+        # Get the algorithm from the token header
+        header = jwt.get_unverified_header(id_token)
+        algorithm = header.get("alg", "RS256")
+
+        # Determine hash algorithm based on JWT algorithm
+        if algorithm.startswith("RS") or algorithm.startswith("PS"):
+            hash_alg = f"sha{algorithm[2:]}"
+        elif algorithm.startswith("ES"):
+            hash_alg = f"sha{algorithm[2:]}"
+        elif algorithm.startswith("HS"):
+            hash_alg = f"sha{algorithm[2:]}"
+        else:
+            hash_alg = "sha256"
+
+        # Calculate expected at_hash
+        try:
+            hash_func = hashlib.new(hash_alg)
+            hash_func.update(access_token.encode("ascii"))
+            digest = hash_func.digest()
+            expected_at_hash = base64.urlsafe_b64encode(digest[: len(digest) // 2]).rstrip(b"=")
+            expected_at_hash_str = expected_at_hash.decode("ascii")
+        except Exception as e:
+            raise TokenValidationError(f"Failed to calculate at_hash: {e}") from e
+
+        if at_hash != expected_at_hash_str:
+            raise TokenValidationError("at_hash claim does not match access token")
+
+    def _extract_identity(
+        self,
+        claims: dict[str, Any],
+        provider_name: str,
+    ) -> IdentityInfo:
+        """Extract identity information from claims.
+
+        Args:
+            claims: The validated token claims.
+            provider_name: The provider name.
+
+        Returns:
+            The extracted identity information.
+        """
+        # Extract extra claims (non-standard)
+        extra_claims = {
+            k: v for k, v in claims.items() if k not in self.STANDARD_CLAIMS
+        }
+
+        return IdentityInfo(
+            provider_name=provider_name,
+            subject=claims["sub"],
+            email=claims.get("email"),
+            email_verified=claims.get("email_verified", False),
+            name=claims.get("name"),
+            given_name=claims.get("given_name"),
+            family_name=claims.get("family_name"),
+            picture=claims.get("picture"),
+            locale=claims.get("locale"),
+            extra_claims=extra_claims,
+        )
+
+
+async def validate_id_token(
+    id_token: str,
+    provider_manager: ProviderManager,
+    provider_name: str,
+    expected_nonce: str | None = None,
+    verify_at_hash: bool = False,
+    access_token: str | None = None,
+) -> IdentityInfo:
+    """Convenience function to validate an ID token.
+
+    Args:
+        id_token: The raw ID token (JWT).
+        provider_manager: The provider manager.
+        provider_name: The provider name.
+        expected_nonce: The expected nonce value.
+        verify_at_hash: Whether to verify the at_hash claim.
+        access_token: The access token (required if verify_at_hash is True).
+
+    Returns:
+        Extracted identity information.
+    """
+    validator = TokenValidator(provider_manager)
+    return await validator.validate(
+        id_token,
+        provider_name,
+        expected_nonce,
+        verify_at_hash,
+        access_token,
+    )