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

"""Microsoft Teams 渠道双通道探针。

通过 Bot Framework Token 获取 + Microsoft Graph API /me 端点
进行双通道健康检查。
"""

from __future__ import annotations

import time
from typing import TYPE_CHECKING

import aiohttp

from yuxi.channels.models import HealthStatus
from yuxi.utils.logging_config import logger

if TYPE_CHECKING:
    from .send import MessageSender

BOT_LOGIN_URL = "https://login.microsoftonline.com/botframework.com/oauth2/v2.0/token"
GRAPH_LOGIN_URL = "https://login.microsoftonline.com/{tenant}/oauth2/v2.0/token"
GRAPH_ME_URL = "https://graph.microsoft.com/v1.0/me"


class MSTeamsProbe:
    def __init__(
        self,
        app_id: str,
        app_password: str,
        tenant_id: str = "",
        sender: MessageSender | None = None,
    ):
        self._app_id = app_id
        self._app_password = app_password
        self._tenant_id = tenant_id
        self._sender = sender
        self._session: aiohttp.ClientSession | None = None

    async def _ensure_session(self) -> aiohttp.ClientSession:
        if self._session is None or self._session.closed:
            self._session = aiohttp.ClientSession()
        return self._session

    async def close(self) -> None:
        if self._session and not self._session.closed:
            await self._session.close()
            self._session = None

    async def probe(self) -> HealthStatus:
        start = time.monotonic()
        issues: list[str] = []

        bot_token = await self._get_bot_token()
        if not bot_token:
            issues.append("Bot Framework auth failed")

        graph_token = await self._get_graph_token()
        if graph_token:
            graph_ok = await self._probe_graph(graph_token)
            if not graph_ok:
                issues.append("Graph API unreachable or unauthorized")
        else:
            issues.append("Graph token unavailable")

        delegated_ok = await self._probe_delegated_auth()
        if not delegated_ok:
            issues.append("Delegated Auth token unavailable")

        latency_ms = (time.monotonic() - start) * 1000

        if issues:
            return HealthStatus(
                status="degraded",
                latency_ms=latency_ms,
                last_error="; ".join(issues),
                metadata={"issues": issues},
            )

        return HealthStatus(
            status="healthy",
            latency_ms=latency_ms,
            metadata={"app_id": self._app_id[:8] + "..."},
        )

    async def _get_bot_token(self) -> str | None:
        if self._sender and self._sender.token:
            return self._sender.token

        data = {
            "client_id": self._app_id,
            "client_secret": self._app_password,
            "grant_type": "client_credentials",
            "scope": "https://api.botframework.com/.default",
        }
        try:
            session = await self._ensure_session()
            async with session.post(BOT_LOGIN_URL, data=data) as resp:
                if resp.status == 200:
                    result = await resp.json()
                    return result.get("access_token")
                logger.warning(f"Bot token request failed: HTTP {resp.status}")
        except Exception as e:
            logger.warning(f"Bot token request error: {e}")
        return None

    async def _get_graph_token(self) -> str | None:
        tenant = self._tenant_id or "common"
        token_url = GRAPH_LOGIN_URL.format(tenant=tenant)
        data = {
            "client_id": self._app_id,
            "client_secret": self._app_password,
            "scope": "https://graph.microsoft.com/.default",
            "grant_type": "client_credentials",
        }
        try:
            session = await self._ensure_session()
            async with session.post(token_url, data=data) as resp:
                if resp.status == 200:
                    result = await resp.json()
                    return result.get("access_token")
        except Exception as e:
            logger.warning(f"Graph token request error: {e}")
        return None

    async def _probe_graph(self, token: str) -> bool:
        try:
            headers = {"Authorization": f"Bearer {token}"}
            session = await self._ensure_session()
            async with session.get(GRAPH_ME_URL, headers=headers) as resp:
                return resp.status == 200
        except Exception as e:
            logger.warning(f"Graph probe error: {e}")
            return False

    async def _probe_delegated_auth(self) -> bool:
        from .credentials import DelegatedAuthStore

        store = DelegatedAuthStore()
        tokens = store._tokens
        if not tokens:
            return False

        for user_id, entry in list(tokens.items())[:1]:
            if entry.get("access_token"):
                headers = {"Authorization": f"Bearer {entry['access_token']}"}
                try:
                    session = await self._ensure_session()
                    async with session.get(GRAPH_ME_URL, headers=headers) as resp:
                        return resp.status == 200
                except Exception as e:
                    logger.warning(f"Delegated auth probe error: {e}")
        return False

    async def validate_credentials(self) -> bool:
        token = await self._get_bot_token()
        return token is not None