import asyncio import time import uuid from datetime import datetime, timedelta, timezone import httpx from sqlalchemy import delete, update from sqlalchemy.ext.asyncio import AsyncSession from app.config import settings from app.models.wecom_bind_token import WecomBindToken async def store_bind_token(db: AsyncSession, wecom_userid: str, ttl: int = 600) -> str: """Persist a one-time bind token so it works across application workers.""" token = uuid.uuid4().hex now = datetime.now(timezone.utc) await db.execute(delete(WecomBindToken).where(WecomBindToken.expires_at < now)) db.add(WecomBindToken( token=token, wecom_userid=wecom_userid, expires_at=now + timedelta(seconds=ttl), )) await db.commit() return token async def consume_bind_token(db: AsyncSession, token: str) -> str | None: """Atomically consume an unexpired binding token and return its WeCom user ID.""" now = datetime.now(timezone.utc) result = await db.execute( update(WecomBindToken) .where( WecomBindToken.token == token, WecomBindToken.consumed_at.is_(None), WecomBindToken.expires_at > now, ) .values(consumed_at=now) .returning(WecomBindToken.wecom_userid) ) await db.commit() return result.scalar_one_or_none() class WecomClient: """WeChat Work API client — token management, message sending, OAuth.""" def __init__(self): self.corp_id = settings.WECOM_CORP_ID self.agent_id = settings.WECOM_AGENT_ID self.secret = settings.WECOM_SECRET self._access_token: str | None = None self._token_expires_at: float = 0 # epoch seconds @property def access_token(self) -> str | None: return self._access_token async def _get_token(self) -> str: """Get a valid access token, refreshing if expired.""" if self._access_token and time.time() < self._token_expires_at - 60: return self._access_token url = f"{settings.WECOM_API_BASE}/cgi-bin/gettoken?corpid={self.corp_id}&corpsecret={self.secret}" async with httpx.AsyncClient() as client: resp = await client.get(url, timeout=10) data = resp.json() if data.get("errcode") == 0: self._access_token = data["access_token"] self._token_expires_at = time.time() + data.get("expires_in", 7200) return self._access_token raise Exception(f"Failed to get wecom token: {data}") async def _post_with_retry(self, url: str, body: dict, max_retries: int = 3) -> dict: """POST with retry on network errors and automatic token refresh on 42001.""" last_error = None for attempt in range(max_retries): try: token = await self._get_token() sep = "&" if "?" in url else "?" full_url = f"{url}{sep}access_token={token}" async with httpx.AsyncClient() as client: resp = await client.post(full_url, json=body, timeout=10) data = resp.json() errcode = data.get("errcode", 0) if errcode == 0: return data if errcode == 42001: self._access_token = None self._token_expires_at = 0 if attempt < max_retries - 1: continue if errcode == 45009: if attempt < max_retries - 1: await asyncio.sleep(1 * (attempt + 1)) continue last_error = data except (httpx.TimeoutException, httpx.ConnectError) as e: last_error = {"errcode": -1, "errmsg": str(e)} if attempt < max_retries - 1: await asyncio.sleep(0.5 * (attempt + 1)) raise Exception(f"WeCom API error after {max_retries} retries: {last_error}") # ── OAuth ── async def get_userinfo_by_code(self, code: str) -> dict | None: """Exchange OAuth2 code for userid (used in silent login).""" try: token = await self._get_token() url = f"{settings.WECOM_API_BASE}/cgi-bin/auth/getuserinfo?access_token={token}&code={code}" async with httpx.AsyncClient() as client: resp = await client.get(url, timeout=10) data = resp.json() if data.get("errcode") == 0: return data return None except Exception: return None # ── Messaging ── async def send_text_message(self, user_ids: list[str], content: str) -> bool: """Send a text app message to specified users.""" if not settings.WECOM_AGENT_ID: return False try: body = { "touser": "|".join(user_ids), "msgtype": "text", "agentid": int(settings.WECOM_AGENT_ID), "text": {"content": content}, } await self._post_with_retry( f"{settings.WECOM_API_BASE}/cgi-bin/message/send", body ) return True except Exception: return False async def send_markdown_message(self, content: str) -> bool: """Send a markdown message to all users in the app (broadcast).""" if not settings.WECOM_AGENT_ID: return False try: body = { "touser": "@all", "msgtype": "markdown", "agentid": int(settings.WECOM_AGENT_ID), "markdown": {"content": content}, } await self._post_with_retry( f"{settings.WECOM_API_BASE}/cgi-bin/message/send", body ) return True except Exception: return False async def send_text_card(self, user_id: str, title: str, description: str, url: str) -> bool: """Send a textcard message (clickable card) to a single user.""" if not settings.WECOM_AGENT_ID: return False try: body = { "touser": user_id, "msgtype": "textcard", "agentid": int(settings.WECOM_AGENT_ID), "textcard": { "title": title, "description": description, "url": url, }, } await self._post_with_retry( f"{settings.WECOM_API_BASE}/cgi-bin/message/send", body ) return True except Exception: return False async def send_template_card( self, user_ids: list[str], title: str, desc: str, url: str, btn_text: str = "查看详情" ) -> bool: """Send a text_notice template card with a deep-link button.""" if not settings.WECOM_AGENT_ID: return False try: body = { "touser": "|".join(user_ids), "msgtype": "template_card", "agentid": int(settings.WECOM_AGENT_ID), "template_card": { "card_type": "text_notice", "main_title": {"title": title, "desc": desc}, "card_action": {"type": 1, "url": url}, "button_list": [{"text": btn_text, "style": 1, "key": "open_url"}], }, } await self._post_with_retry( f"{settings.WECOM_API_BASE}/cgi-bin/message/send", body ) return True except Exception: return False # ── Menu management ── async def create_menu(self, buttons: list[dict]) -> bool: """Create/replace the app's custom menu.""" if not settings.WECOM_AGENT_ID: return False try: body = {"button": buttons} await self._post_with_retry( f"{settings.WECOM_API_BASE}/cgi-bin/menu/create?agentid={settings.WECOM_AGENT_ID}", body ) return True except Exception: return False async def get_menu(self) -> dict | None: """Get current app menu configuration.""" if not settings.WECOM_AGENT_ID: return None try: token = await self._get_token() url = f"{settings.WECOM_API_BASE}/cgi-bin/menu/get?access_token={token}&agentid={settings.WECOM_AGENT_ID}" async with httpx.AsyncClient() as client: resp = await client.get(url, timeout=10) data = resp.json() if data.get("errcode") == 0: return data return None except Exception: return None async def delete_menu(self) -> bool: """Delete the app's custom menu.""" if not settings.WECOM_AGENT_ID: return False try: token = await self._get_token() url = f"{settings.WECOM_API_BASE}/cgi-bin/menu/delete?access_token={token}&agentid={settings.WECOM_AGENT_ID}" async with httpx.AsyncClient() as client: resp = await client.get(url, timeout=10) data = resp.json() return data.get("errcode") == 0 except Exception: return False wecom_client = WecomClient()