ForcePilot/backend/package/yuxi/channel/extensions/msteams/jwks.py

from __future__ import annotations

import logging
from typing import Any

import httpx
from cachetools import TTLCache
from cryptography.hazmat.primitives import serialization

logger = logging.getLogger(__name__)

JWKS_CACHE_MAXSIZE = 100
JWKS_CACHE_TTL = 600

_jwks_cache: TTLCache = TTLCache(maxsize=JWKS_CACHE_MAXSIZE, ttl=JWKS_CACHE_TTL)


class JWKSError(Exception):
    pass


class JWKSNetworkError(JWKSError):
    pass


class JWKSKeyNotFoundError(JWKSError):
    pass


ISSUER_JWKS_MAP = {
    "https://api.botframework.com": "https://login.botframework.com/v1/.well-known/keys",
    "login.microsoftonline.com": "https://login.microsoftonline.com/common/discovery/v2.0/keys",
    "sts.windows.net": "https://login.microsoftonline.com/common/discovery/v2.0/keys",
}


def _resolve_jwks_uri(iss: str) -> str | None:
    for prefix, uri in ISSUER_JWKS_MAP.items():
        if prefix in iss or iss.startswith(prefix):
            return uri
    if "login.microsoftonline.com" in iss:
        return ISSUER_JWKS_MAP["login.microsoftonline.com"]
    if "sts.windows.net" in iss:
        return ISSUER_JWKS_MAP["sts.windows.net"]
    return None


def _pem_from_jwk(jwk: dict) -> str:
    from cryptography.hazmat.backends import default_backend

    kty = jwk.get("kty")

    if kty == "RSA":
        from cryptography.hazmat.primitives.asymmetric.rsa import RSAPublicNumbers

        n = int.from_bytes(_b64_decode(_ensure_str(jwk["n"])), byteorder="big")
        e = int.from_bytes(_b64_decode(_ensure_str(jwk["e"])), byteorder="big")
        pub = RSAPublicNumbers(e, n).public_key(default_backend())
        return pub.public_bytes(serialization.Encoding.PEM, serialization.PublicFormat.SubjectPublicKeyInfo).decode(
            "ascii"
        )

    if kty == "EC":
        from cryptography.hazmat.primitives.asymmetric.ec import EllipticCurvePublicNumbers

        crv = jwk.get("crv", "P-256")
        x = int.from_bytes(_b64_decode(_ensure_str(jwk["x"])), byteorder="big")
        y = int.from_bytes(_b64_decode(_ensure_str(jwk["y"])), byteorder="big")

        import cryptography.hazmat.primitives.asymmetric.ec as _ec

        curve = _ec.SECP256R1() if crv == "P-256" else _ec.SECP384R1()
        pub_numbers = EllipticCurvePublicNumbers(x, y, curve)
        pub = pub_numbers.public_key(default_backend())
        return pub.public_bytes(serialization.Encoding.PEM, serialization.PublicFormat.SubjectPublicKeyInfo).decode(
            "ascii"
        )

    raise JWKSKeyNotFoundError(f"Unsupported key type: {kty}")


def _b64_decode(data: str) -> bytes:
    import base64

    padded = data + "=" * (-len(data) % 4)
    return base64.urlsafe_b64decode(padded)


def _ensure_str(val: Any) -> str:
    if isinstance(val, str):
        return val
    return str(val)


async def _fetch_jwks(jwks_uri: str) -> dict:
    cache_key = f"jwks:{jwks_uri}"
    cached = _jwks_cache.get(cache_key)
    if cached:
        return cached

    try:
        async with httpx.AsyncClient(timeout=15.0) as client:
            resp = await client.get(jwks_uri)
            resp.raise_for_status()
            data = resp.json()
            _jwks_cache[cache_key] = data
            return data
    except httpx.NetworkError as e:
        raise JWKSNetworkError(f"Network error fetching JWKS from {jwks_uri}: {e}") from e
    except httpx.HTTPStatusError as e:
        raise JWKSNetworkError(f"HTTP {e.response.status_code} fetching JWKS from {jwks_uri}") from e
    except Exception as e:
        raise JWKSError(f"Failed to fetch JWKS from {jwks_uri}: {e}") from e


async def get_jwks_public_key(jwks_uri: str, kid: str) -> str:
    jwks_data = await _fetch_jwks(jwks_uri)
    keys = jwks_data.get("keys", [])

    for jwk in keys:
        if jwk.get("kid") == kid:
            return _pem_from_jwk(jwk)

    raise JWKSKeyNotFoundError(f"Key with kid '{kid}' not found in JWKS from {jwks_uri}")


def collect_all_issuers(tenant_id: str) -> list[str]:
    return [
        "https://api.botframework.com",
        f"https://login.microsoftonline.com/{tenant_id}/v2.0",
        f"https://sts.windows.net/{tenant_id}/",
    ]


def resolve_jwks_uri_for_iss(iss: str) -> str | None:
    return _resolve_jwks_uri(iss)


def clear_jwks_cache() -> None:
    _jwks_cache.clear()