ForcePilot/backend/package/yuxi/channels/adapters/feishu/oauth.py

from __future__ import annotations

import time
from typing import Any

from yuxi.utils.logging_config import logger

TOKEN_CACHE_TTL_S = 3600
TOKEN_REFRESH_MARGIN_S = 300


DEFAULT_FEISHU_DOMAIN = "https://open.feishu.cn"
DEFAULT_LARK_DOMAIN = "https://open.larksuite.com"


class FeishuOAuthClient:
    def __init__(self, app_id: str, app_secret: str, redirect_uri: str = "", domain: str = ""):
        self._app_id = app_id
        self._app_secret = app_secret
        self._redirect_uri = redirect_uri
        self._domain = domain or DEFAULT_FEISHU_DOMAIN
        self._token_cache: dict[str, tuple[dict[str, Any], float]] = {}
        self._access_token: str = ""
        self._access_token_expire_at: float = 0
        self._refresh_token: str = ""

    def get_authorization_url(self, state: str = "", scope: str = "") -> str:
        base_url = f"{self._domain}/open-apis/authen/v1/index"
        params = {
            "app_id": self._app_id,
            "redirect_uri": self._redirect_uri or "http://localhost/callback",
        }
        if state:
            params["state"] = state
        if scope:
            params["scope"] = scope
        query = "&".join(f"{k}={v}" for k, v in params.items())
        return f"{base_url}?{query}"

    async def exchange_code_for_token(self, code: str) -> dict[str, Any]:
        import httpx

        cache_key = f"code:{code}"
        cached = self._token_cache.get(cache_key)
        if cached is not None:
            token_data, cached_at = cached
            if time.monotonic() - cached_at < 30:
                return token_data
            del self._token_cache[cache_key]

        url = f"{self._domain}/open-apis/authen/v1/oidc/access_token"
        headers = {"Content-Type": "application/json"}
        body = {
            "app_id": self._app_id,
            "app_secret": self._app_secret,
            "grant_type": "authorization_code",
            "code": code,
        }

        async with httpx.AsyncClient(timeout=httpx.Timeout(30)) as client:
            resp = await client.post(url, headers=headers, json=body)
            if resp.status_code == 200:
                data = resp.json()
                self._token_cache[cache_key] = (data, time.monotonic())
                if data.get("access_token"):
                    self._access_token = data["access_token"]
                    expires_in = data.get("expires_in", 7200)
                    self._access_token_expire_at = time.monotonic() + expires_in - TOKEN_REFRESH_MARGIN_S
                if data.get("refresh_token"):
                    self._refresh_token = data["refresh_token"]
                return data
            logger.error("[FeishuOAuth] Token exchange failed: HTTP %d", resp.status_code)
            return {}

    async def refresh_user_access_token(self, refresh_token: str = "") -> dict[str, Any]:
        import httpx

        rt = refresh_token or self._refresh_token
        if not rt:
            logger.error("[FeishuOAuth] No refresh token available")
            return {}

        cache_key = f"refresh:{rt}"
        cached = self._token_cache.get(cache_key)
        if cached is not None:
            token_data, cached_at = cached
            if time.monotonic() - cached_at < 30:
                return token_data
            del self._token_cache[cache_key]

        url = f"{self._domain}/open-apis/authen/v1/oidc/refresh_access_token"
        headers = {"Content-Type": "application/json"}
        body = {
            "app_id": self._app_id,
            "app_secret": self._app_secret,
            "grant_type": "refresh_token",
            "refresh_token": rt,
        }

        async with httpx.AsyncClient(timeout=httpx.Timeout(30)) as client:
            resp = await client.post(url, headers=headers, json=body)
            if resp.status_code == 200:
                data = resp.json()
                self._token_cache[cache_key] = (data, time.monotonic())
                if data.get("access_token"):
                    self._access_token = data["access_token"]
                    expires_in = data.get("expires_in", 7200)
                    self._access_token_expire_at = time.monotonic() + expires_in - TOKEN_REFRESH_MARGIN_S
                if data.get("refresh_token"):
                    self._refresh_token = data["refresh_token"]
                return data
            logger.error("[FeishuOAuth] Refresh failed: HTTP %d", resp.status_code)
            return {}

    def is_access_token_expired(self) -> bool:
        if not self._access_token:
            return True
        return time.monotonic() >= self._access_token_expire_at

    def clear_cache(self) -> None:
        self._token_cache.clear()
        self._access_token = ""
        self._access_token_expire_at = 0
        self._refresh_token = ""


class FeishuDeviceCodeClient:

    def __init__(self, app_id: str, app_secret: str, domain: str = ""):
        self._app_id = app_id
        self._app_secret = app_secret
        self._domain = domain or DEFAULT_FEISHU_DOMAIN
        self._device_code_url = f"{self._domain}/open-apis/authen/v1/device/code"
        self._device_token_url = f"{self._domain}/open-apis/authen/v1/oidc/access_token"
        self._token_cache: dict[str, tuple[dict[str, Any], float]] = {}

    def init_device_flow(self) -> dict[str, Any] | None:
        import httpx

        body = {
            "app_id": self._app_id,
            "app_secret": self._app_secret,
            "scope": "user:read",
        }
        try:
            resp = httpx.post(
                self._device_code_url,
                json=body,
                headers={"Content-Type": "application/json"},
                timeout=httpx.Timeout(30),
            )
            if resp.status_code == 200:
                data = resp.json()
                if data.get("code") == 0:
                    return data.get("data", {})
                logger.error("[FeishuDeviceCode] init failed: code=%s, msg=%s", data.get("code"), data.get("msg"))
            else:
                logger.error("[FeishuDeviceCode] init HTTP %d: %s", resp.status_code, resp.text[:300])
        except Exception as e:
            logger.error("[FeishuDeviceCode] init error: %s", e)
        return None

    def poll_device_token(self, device_code: str) -> dict[str, Any] | None:
        import httpx

        cache_key = f"device:{device_code}"
        cached = self._token_cache.get(cache_key)
        if cached is not None:
            token_data, cached_at = cached
            if time.monotonic() - cached_at < 5:
                return token_data
            del self._token_cache[cache_key]

        body = {
            "app_id": self._app_id,
            "app_secret": self._app_secret,
            "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
            "device_code": device_code,
        }
        try:
            resp = httpx.post(
                self._device_token_url,
                json=body,
                headers={"Content-Type": "application/json"},
                timeout=httpx.Timeout(30),
            )
            if resp.status_code == 200:
                data = resp.json()
                self._token_cache[cache_key] = (data, time.monotonic())
                return data
            logger.error("[FeishuDeviceCode] poll HTTP %d: %s", resp.status_code, resp.text[:300])
        except Exception as e:
            logger.error("[FeishuDeviceCode] poll error: %s", e)
        return None

    def clear_cache(self) -> None:
        self._token_cache.clear()