from __future__ import annotations import asyncio import json import logging import time from typing import TYPE_CHECKING import httpx from yuxi.channel.extensions.dingtalk.types import DingTalkAccount if TYPE_CHECKING: from yuxi.channel.extensions.dingtalk.monitor import DingTalkMonitor logger = logging.getLogger(__name__) TOKEN_URL = "https://api.dingtalk.com/v1.0/oauth2/accessToken" TOKEN_REFRESH_MARGIN = 300 RECONNECT_WAIT_MS = 100 RECONNECT_MAX_WAITS = 100 OPEN_CONNECTION_URL = "https://api.dingtalk.com/v1.0/gateway/connections/open" class DingTalkTokenManager: def __init__(self, app_key: str, app_secret: str): self._app_key = app_key self._app_secret = app_secret self._token: str | None = None self._expires_at: float = 0.0 self._lock = asyncio.Lock() self._http: httpx.AsyncClient | None = None def set_http(self, http: httpx.AsyncClient) -> None: self._http = http @property def token(self) -> str | None: if self._token and time.time() < self._expires_at: return self._token return None async def get_access_token(self) -> str | None: if self._token and time.time() < self._expires_at: return self._token async with self._lock: if self._token and time.time() < self._expires_at: return self._token try: resp = await self._http.post( TOKEN_URL, json={ "appKey": self._app_key, "appSecret": self._app_secret, }, timeout=10.0, ) data = resp.json() self._token = data["accessToken"] expire_in = data.get("expireIn", 7200) self._expires_at = time.time() + expire_in - TOKEN_REFRESH_MARGIN logger.info("DingTalk access_token refreshed, expires_in=%ds", expire_in) return self._token except Exception: logger.exception("Failed to get DingTalk access_token") return None class DingTalkGateway: def __init__(self, account: DingTalkAccount, token_manager: DingTalkTokenManager): self._account = account self._token_manager = token_manager self._running = False self._event_loop: asyncio.AbstractEventLoop | None = None self._robot_code_cache: str | None = None self._handler: DingTalkMonitor | None = None self._stream_client = None self._credential = None def set_handler(self, handler: DingTalkMonitor) -> None: self._handler = handler @property def robot_code(self) -> str | None: return self._robot_code_cache or self._account.robot_code @property def app_key(self) -> str: return self._account.app_key @property def app_secret(self) -> str: return self._account.app_secret @property def token_manager(self) -> DingTalkTokenManager: return self._token_manager async def start(self, ctx) -> object: if not self._account.is_configured(): logger.warning("DingTalk account not configured, skipping start") return {"running": False, "reason": "not-configured"} self._running = True self._event_loop = asyncio.get_running_loop() try: from dingtalk_stream import ChatbotMessage, Credential, DingTalkStreamClient self._credential = Credential(self._account.app_key, self._account.app_secret) self._stream_client = DingTalkStreamClient( credential=self._credential, ) self._stream_client.register_callback_handler(ChatbotMessage.TOPIC, self._handler) extra_topics = getattr(self._account, "subscribe_topics", []) for topic in extra_topics: try: self._stream_client.register_callback_handler(topic, self._handler) logger.info("DingTalk extra topic registered: %s", topic) except Exception: logger.warning("Failed to register topic: %s", topic, exc_info=True) try: self._stream_client.register_callback_handler("/v1.0/card/instances/callback", self._handler) logger.info("DingTalk card callback registered") except Exception: logger.debug("Failed to register card callback", exc_info=True) self._stream_client.pre_start() except ImportError: logger.error("dingtalk-stream-sdk-python not installed") return {"running": False, "reason": "sdk-missing"} except Exception: logger.exception("Failed to initialize DingTalk stream client") return {"running": False, "reason": "init-failed"} asyncio.create_task(self._connection_loop(), name="dingtalk-stream-loop") logger.info("DingTalk gateway started for account %s", self._account.account_id) return {"running": True, "account_id": self._account.account_id} async def stop(self, ctx) -> None: self._running = False self._stream_client = None logger.info("DingTalk gateway stopped") async def _open_connection(self) -> tuple[str, str]: async with httpx.AsyncClient(timeout=10.0) as http: resp = await http.post( OPEN_CONNECTION_URL, headers={ "Content-Type": "application/json", "clientId": self._account.app_key, "dingtalk-accept-encoding": "base64", }, ) resp.raise_for_status() data = resp.json() return data["endpoint"], data["ticket"] async def _connection_loop(self): first_attempt = True while self._running: try: endpoint, ticket = await self._open_connection() except Exception as e: if first_attempt: logger.error("DingTalk first connection failed: %s", e) first_attempt = False await self._wait_with_check(RECONNECT_MAX_WAITS) continue try: import websockets async with websockets.connect(f"{endpoint}?ticket={ticket}") as ws: if first_attempt: logger.info("DingTalk stream connected") first_attempt = False async for raw_message in ws: if not self._running: break try: json_message = json.loads(raw_message) result = self._stream_client.route_message(json_message) if result == "TAG_DISCONNECT": logger.warning("DingTalk stream: TAG_DISCONNECT received") break except Exception: logger.exception("DingTalk stream message routing error") except Exception: logger.exception("DingTalk stream connection error") await self._wait_with_check(RECONNECT_MAX_WAITS) async def _wait_with_check(self, max_waits: int): for _ in range(max_waits): if not self._running: break await asyncio.sleep(RECONNECT_WAIT_MS / 1000) async def probe(self, account: dict) -> bool: return account.get("app_key") and account.get("app_secret")