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

from __future__ import annotations

import logging
import secrets
import time
from dataclasses import dataclass, field

logger = logging.getLogger(__name__)

SSO_STATE_TTL_SECONDS = 300


@dataclass
class SSOTokenState:
    state: str
    token: str = ""
    created_at: float = field(default_factory=time.monotonic)
    exchanged: bool = False
    user_id: str = ""
    tenant_id: str = ""
    error: str = ""

    @property
    def is_expired(self) -> bool:
        return time.monotonic() - self.created_at > SSO_STATE_TTL_SECONDS


class SSOTokenStore:
    def __init__(self):
        self._states: dict[str, SSOTokenState] = {}

    def create_state(self, tenant_id: str = "", user_id: str = "") -> str:
        state = secrets.token_hex(16)
        self._states[state] = SSOTokenState(state=state, tenant_id=tenant_id, user_id=user_id)
        self._cleanup_expired()
        return state

    def verify_state(self, state: str) -> SSOTokenState | None:
        stored = self._states.get(state)
        if not stored:
            return None
        if stored.is_expired:
            self._states.pop(state, None)
            return None
        return stored

    def exchange_token(self, state: str, token: str) -> bool:
        stored = self.verify_state(state)
        if not stored:
            return False
        stored.token = token
        stored.exchanged = True
        return True

    def revoke_state(self, state: str) -> None:
        self._states.pop(state, None)

    def list_active_states(self) -> list[dict]:
        self._cleanup_expired()
        return [
            {
                "state": s.state,
                "exchanged": s.exchanged,
                "user_id": s.user_id,
                "tenant_id": s.tenant_id,
            }
            for s in self._states.values()
        ]

    def _cleanup_expired(self) -> None:
        expired = [s for s, st in self._states.items() if st.is_expired]
        for s in expired:
            self._states.pop(s, None)


async def exchange_token_with_microsoft(
    tenant_id: str,
    client_id: str,
    client_secret: str,
    code: str,
    redirect_uri: str,
) -> dict | None:
    try:
        import httpx

        token_url = f"https://login.microsoftonline.com/{tenant_id}/oauth2/v2.0/token"
        body = {
            "client_id": client_id,
            "client_secret": client_secret,
            "code": code,
            "redirect_uri": redirect_uri,
            "grant_type": "authorization_code",
        }
        async with httpx.AsyncClient(timeout=30.0) as client:
            resp = await client.post(token_url, data=body)
            resp.raise_for_status()
            return resp.json()
    except Exception as e:
        logger.warning("Token exchange failed: %s", e)
        return None


async def exchange_token_on_behalf_of(
    tenant_id: str,
    client_id: str,
    client_secret: str,
    user_token: str,
    scope: str,
) -> dict | None:
    try:
        import httpx

        token_url = f"https://login.microsoftonline.com/{tenant_id}/oauth2/v2.0/token"
        body = {
            "client_id": client_id,
            "client_secret": client_secret,
            "assertion": user_token,
            "scope": scope,
            "grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer",
            "requested_token_use": "on_behalf_of",
        }
        async with httpx.AsyncClient(timeout=30.0) as client:
            resp = await client.post(token_url, data=body)
            resp.raise_for_status()
            return resp.json()
    except Exception as e:
        logger.warning("OBO token exchange failed: %s", e)
        return None


def build_sso_card(connection_name: str, state: str, text: str = "") -> dict:
    card = {
        "type": "AdaptiveCard",
        "version": "1.4",
        "body": [
            {
                "type": "TextBlock",
                "text": text or "🔐 登录认证",
                "weight": "Bolder",
                "size": "Medium",
            },
            {
                "type": "TextBlock",
                "text": "请点击下方按钮完成 Microsoft 365 身份认证",
                "wrap": True,
            },
        ],
        "actions": [
            {
                "type": "Action.Submit",
                "title": "🔑 登录 Microsoft 365",
                "data": {
                    "action": "sso_signin",
                    "state": state,
                    "connection_name": connection_name,
                },
            },
        ],
    }
    return card


def build_sso_attachment(connection_name: str, state: str, text: str = "") -> dict:
    from .adaptive_card import build_attachment

    card = build_sso_card(connection_name, state, text)
    return build_attachment(card)