253 lines
9.3 KiB
Python
253 lines
9.3 KiB
Python
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:
|
|
time.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:
|
|
time.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()
|