80 lines
2.7 KiB
Python
80 lines
2.7 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import time
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
from yuxi.utils.logging_config import logger
|
|
|
|
TOKEN_REFRESH_MARGIN_S = 200
|
|
TOKEN_VALIDITY_S = 7200
|
|
|
|
|
|
class DingDingTokenManager:
|
|
def __init__(self, app_key: str, app_secret: str, http_client: httpx.AsyncClient | None = None):
|
|
self._app_key = app_key
|
|
self._app_secret = app_secret
|
|
self._access_token: str | None = None
|
|
self._token_expires_at: float = 0
|
|
self._lock = asyncio.Lock()
|
|
self._http_client = http_client
|
|
|
|
async def get_token(self) -> str:
|
|
if self._is_valid():
|
|
return self._access_token # type: ignore[return-value]
|
|
|
|
async with self._lock:
|
|
if self._is_valid():
|
|
return self._access_token # type: ignore[return-value]
|
|
await self._refresh()
|
|
|
|
if not self._access_token:
|
|
raise RuntimeError("Failed to obtain DingDing access token")
|
|
return self._access_token
|
|
|
|
async def _refresh(self) -> None:
|
|
refresh_start = time.monotonic()
|
|
url = "https://api.dingtalk.com/v1.0/oauth2/accessToken"
|
|
payload: dict[str, Any] = {
|
|
"appKey": self._app_key,
|
|
"appSecret": self._app_secret,
|
|
}
|
|
|
|
if self._http_client is not None:
|
|
resp = await self._http_client.post(url, json=payload)
|
|
else:
|
|
async with httpx.AsyncClient(timeout=httpx.Timeout(15)) as client:
|
|
resp = await client.post(url, json=payload)
|
|
|
|
if resp.status_code != 200:
|
|
logger.error(f"[DingDing] Token refresh HTTP {resp.status_code}: {resp.text[:200]}")
|
|
raise RuntimeError(f"Token refresh failed: HTTP {resp.status_code}")
|
|
|
|
data = resp.json()
|
|
token = data.get("accessToken")
|
|
if not token:
|
|
logger.error(f"[DingDing] Token refresh response missing accessToken: {data}")
|
|
raise RuntimeError("Token refresh response missing accessToken")
|
|
|
|
expire_in = data.get("expireIn", TOKEN_VALIDITY_S)
|
|
self._access_token = token
|
|
self._token_expires_at = time.time() + expire_in - TOKEN_REFRESH_MARGIN_S
|
|
|
|
elapsed = time.monotonic() - refresh_start
|
|
if elapsed > 2.0:
|
|
logger.warning(f"[DingDing] Token refresh took {elapsed:.1f}s")
|
|
logger.info(f"[DingDing] Token refreshed, expires in {expire_in - TOKEN_REFRESH_MARGIN_S}s")
|
|
|
|
def _is_valid(self) -> bool:
|
|
return self._access_token is not None and time.time() < self._token_expires_at
|
|
|
|
def invalidate(self) -> None:
|
|
self._access_token = None
|
|
self._token_expires_at = 0
|
|
|
|
async def close(self) -> None:
|
|
self.invalidate()
|
|
self._http_client = None
|