TorbatYar/backend/shared-lib/shared/auth/jwt.py

"""اعتبارسنجی JWT مشترک (Keycloak / OIDC) برای همه سرویس‌های backend."""
from __future__ import annotations

import time
from dataclasses import dataclass
from typing import Any

import httpx
import jwt

from shared.exceptions import UnauthorizedError
from shared.security import CurrentUser


@dataclass(frozen=True)
class JWTSettings:
    keycloak_enabled: bool = False
    keycloak_server_url: str = "http://localhost:8080"
    keycloak_realm: str = "superapp"
    jwt_algorithm: str = "RS256"
    jwt_audience: str = "account"
    jwt_verify_signature: bool = False
    # issuer واقعی داخل JWT (ممکن است با keycloak_server_url داخلی Docker متفاوت باشد)
    jwt_issuer: str | None = None

    @property
    def jwks_url(self) -> str:
        return (
            f"{self.keycloak_server_url}/realms/{self.keycloak_realm}"
            "/protocol/openid-connect/certs"
        )

    @property
    def issuer(self) -> str:
        if self.jwt_issuer:
            return self.jwt_issuer
        return f"{self.keycloak_server_url}/realms/{self.keycloak_realm}"

    @property
    def allowed_issuers(self) -> list[str]:
        """Issuerهای مجاز — توکن ممکن است از URL داخلی Docker یا عمومی صادر شده باشد."""
        issuers = [
            self.issuer,
            f"{self.keycloak_server_url}/realms/{self.keycloak_realm}",
        ]
        if self.jwt_issuer:
            issuers.insert(0, self.jwt_issuer)
        return list(dict.fromkeys(i for i in issuers if i))


class JWTValidator:
    """اعتبارسنج JWT با پشتیبانی JWKS و حالت development."""

    def __init__(self, settings: JWTSettings) -> None:
        self.settings = settings
        self._jwks_cache: dict[str, Any] = {"keys": None, "fetched_at": 0.0}
        self._jwks_ttl = 3600

    async def validate(self, token: str) -> CurrentUser:
        if not token:
            raise UnauthorizedError("توکن احراز هویت ارائه نشده است")
        try:
            if self.settings.jwt_verify_signature and self.settings.keycloak_enabled:
                claims = await self._decode_verified(token)
            else:
                claims = self._decode_unverified(token)
        except UnauthorizedError:
            raise
        except Exception as exc:
            raise UnauthorizedError(f"توکن نامعتبر است: {exc}") from exc
        return self._claims_to_user(claims)

    def _decode_unverified(self, token: str) -> dict[str, Any]:
        return jwt.decode(token, options={"verify_signature": False})

    async def _decode_verified(self, token: str) -> dict[str, Any]:
        jwks = await self._get_jwks()
        if jwks is None:
            raise UnauthorizedError("امکان دریافت کلید عمومی برای اعتبارسنجی توکن نبود")
        try:
            signing_key = jwt.PyJWKClient(self.settings.jwks_url).get_signing_key_from_jwt(token)
            claims = jwt.decode(
                token,
                signing_key.key,
                algorithms=[self.settings.jwt_algorithm],
                options={"verify_aud": False, "verify_iss": False},
            )
            token_iss = claims.get("iss")
            if not self._issuer_allowed(token_iss):
                raise UnauthorizedError(
                    f"توکن نامعتبر است: issuer نامعتبر ({token_iss})"
                )
            return claims
        except jwt.PyJWTError as exc:
            raise UnauthorizedError(f"توکن نامعتبر است: {exc}") from exc

    def _issuer_allowed(self, token_iss: str | None) -> bool:
        if not token_iss:
            return False
        if token_iss in self.settings.allowed_issuers:
            return True
        realm_suffix = f"/realms/{self.settings.keycloak_realm}"
        return token_iss.endswith(realm_suffix)

    async def _get_jwks(self) -> dict[str, Any] | None:
        now = time.time()
        if (
            self._jwks_cache["keys"] is not None
            and now - self._jwks_cache["fetched_at"] < self._jwks_ttl
        ):
            return self._jwks_cache["keys"]
        try:
            async with httpx.AsyncClient(timeout=5.0) as client:
                resp = await client.get(self.settings.jwks_url)
                resp.raise_for_status()
                keys = resp.json()
                self._jwks_cache["keys"] = keys
                self._jwks_cache["fetched_at"] = now
                return keys
        except Exception:
            return None

    @staticmethod
    def _claims_to_user(claims: dict[str, Any]) -> CurrentUser:
        realm_access = claims.get("realm_access") or {}
        roles = list(realm_access.get("roles", [])) if isinstance(realm_access, dict) else []
        # نقش‌های client-specific (مثلاً superapp-frontend)
        resource_access = claims.get("resource_access") or {}
        if isinstance(resource_access, dict):
            for client_roles in resource_access.values():
                if isinstance(client_roles, dict):
                    roles.extend(client_roles.get("roles", []))
        roles = list(dict.fromkeys(roles))  # یکتا
        return CurrentUser(
            user_id=str(claims.get("sub", "")),
            username=claims.get("preferred_username"),
            email=claims.get("email"),
            roles=roles,
            tenant_id=claims.get("tenant_id"),
        )