"""Microsoft Teams 凭据管理。 支持 Secret、Certificate (Federated) 和 Managed Identity 三种凭据模式， 以及 Delegated Auth (OAuth 2.0 Refresh Token)。 """ from __future__ import annotations import json import os import time from pathlib import Path from typing import Any import aiohttp from yuxi.utils.logging_config import logger TOKEN_URL_TEMPLATE = "https://login.microsoftonline.com/{tenant}/oauth2/v2.0/token" BOT_SCOPE = "https://api.botframework.com/.default" GRAPH_SCOPE = "https://graph.microsoft.com/.default" TOKEN_GRACE_PERIOD_S = 60 DELEGATED_STORE_FILENAME = "msteams-delegated.json" DELEGATED_REFRESH_BUFFER_S = 300 class FederatedCredential: def __init__( self, client_id: str, tenant_id: str = "common", certificate_path: str = "", certificate_thumbprint: str = "", use_managed_identity: bool = False, ): self.client_id = client_id self.tenant_id = tenant_id self.certificate_path = certificate_path self.certificate_thumbprint = certificate_thumbprint self.use_managed_identity = use_managed_identity @property def auth_type(self) -> str: if self.use_managed_identity: return "managed_identity" if self.certificate_path: return "certificate" return "secret" async def get_token(self, scope: str) -> str | None: if self.use_managed_identity: return await self._get_managed_identity_token(scope) if self.certificate_path: return await self._get_certificate_token(scope) return None async def _get_managed_identity_token(self, scope: str) -> str | None: endpoint = os.getenv("MSI_ENDPOINT", "http://169.254.169.254/metadata/identity/oauth2/token") params = { "api-version": "2019-08-01", "resource": scope.rstrip("/.default"), "client_id": self.client_id, } try: async with aiohttp.ClientSession() as session: async with session.get( endpoint, params=params, headers={"Metadata": "true"}, ) as resp: if resp.status == 200: result = await resp.json() return result.get("access_token") logger.warning(f"Managed Identity token failed: HTTP {resp.status}") except Exception as e: logger.warning(f"Managed Identity token error: {e}") return None async def _get_certificate_token(self, scope: str) -> str | None: try: from cryptography.hazmat.primitives import serialization from cryptography.hazmat.backends import default_backend import jwt as pyjwt cert_path = Path(self.certificate_path) if not cert_path.exists(): logger.error(f"Certificate not found: {self.certificate_path}") return None with open(cert_path, "rb") as f: private_key = serialization.load_pem_private_key(f.read(), password=None, backend=default_backend()) token_url = TOKEN_URL_TEMPLATE.format(tenant=self.tenant_id) now = int(time.time()) assertion = pyjwt.encode( { "aud": token_url, "iss": self.client_id, "sub": self.client_id, "jti": os.urandom(16).hex(), "nbf": now, "exp": now + 600, }, private_key, algorithm="RS256", headers={"x5t": self.certificate_thumbprint} if self.certificate_thumbprint else {}, ) data = { "client_id": self.client_id, "client_assertion": assertion, "client_assertion_type": "urn:ietf:params:oauth:client-assertion-type:jwt-bearer", "scope": scope, "grant_type": "client_credentials", } async with aiohttp.ClientSession() as session: async with session.post(token_url, data=data) as resp: if resp.status == 200: result = await resp.json() return result.get("access_token") body = await resp.text() logger.warning(f"Certificate token failed: HTTP {resp.status} - {body[:200]}") except ImportError: logger.error("cryptography and PyJWT required for certificate auth") except Exception as e: logger.error(f"Certificate token error: {e}") return None class DelegatedAuthStore: def __init__(self, storage_dir: str | None = None): self._storage_dir = Path(storage_dir or str(Path.home() / ".yuxi" / "msteams")) self._storage_dir.mkdir(parents=True, exist_ok=True) self._tokens: dict[str, dict[str, Any]] = {} self._load() @property def file_path(self) -> Path: return self._storage_dir / DELEGATED_STORE_FILENAME def _load(self) -> None: if not self.file_path.exists(): return try: data = json.loads(self.file_path.read_text(encoding="utf-8")) except (json.JSONDecodeError, OSError): return self._tokens = data.get("tokens", {}) def _save(self) -> None: try: self.file_path.write_text( json.dumps({"tokens": self._tokens}, ensure_ascii=False, indent=2), encoding="utf-8", ) except OSError as e: logger.error(f"MSTeams delegated store: failed to save: {e}") def store_token(self, user_id: str, access_token: str, refresh_token: str, expires_in: int = 3600) -> None: self._tokens[user_id] = { "access_token": access_token, "refresh_token": refresh_token, "expires_at": time.time() + expires_in, } self._save() def get_token(self, user_id: str) -> dict[str, Any] | None: entry = self._tokens.get(user_id) if not entry: return None return entry async def refresh_token( self, user_id: str, client_id: str, client_secret: str, ) -> str | None: entry = self._tokens.get(user_id) if not entry or not entry.get("refresh_token"): return None token_url = TOKEN_URL_TEMPLATE.format(tenant="common") data = { "client_id": client_id, "client_secret": client_secret, "refresh_token": entry["refresh_token"], "grant_type": "refresh_token", "scope": GRAPH_SCOPE, } try: async with aiohttp.ClientSession() as session: async with session.post(token_url, data=data) as resp: if resp.status == 200: result = await resp.json() new_token = result.get("access_token", "") new_refresh = result.get("refresh_token", "") if new_token: self.store_token( user_id, new_token, new_refresh or entry["refresh_token"], result.get("expires_in", 3600), ) return new_token logger.warning(f"Delegated token refresh failed: HTTP {resp.status}") except Exception as e: logger.error(f"Delegated token refresh error: {e}") return None def is_expired(self, user_id: str) -> bool: entry = self._tokens.get(user_id) if not entry: return True return time.time() >= entry.get("expires_at", 0) - DELEGATED_REFRESH_BUFFER_S