ForcePilot/backend/package/yuxi/channels/adapters/msteams/credentials.py

"""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