qiji/backend/app/api/wecom.py

import hashlib
import hmac
import struct
import base64
import uuid
from typing import Optional
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from defusedxml import ElementTree as ET
from fastapi import APIRouter, Depends, HTTPException, Query, Request
from fastapi.responses import PlainTextResponse
from pydantic import BaseModel
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select
from sqlalchemy.exc import IntegrityError
from app.config import settings
from app.database import async_session, get_db
from app.middleware.auth import get_current_user, require_director
from app.models.user import User
from app.services.wecom import wecom_client, store_bind_token, consume_bind_token
from app.services.scheduler import check_daily_reporting

router = APIRouter(prefix="/wecom", tags=["WeChatWork"])


class RemindRequest(BaseModel):
    user_ids: list[str]
    message: Optional[str] = None


class AnnouncementRequest(BaseModel):
    content: str


class BindConfirmRequest(BaseModel):
    token: str


# ── Crypto helpers ──

def _verify_signature(token: str, timestamp: str, nonce: str, encrypted: str, signature: str) -> bool:
    items = sorted([token, timestamp, nonce, encrypted])
    raw = "".join(items)
    # The WeCom callback protocol mandates SHA-1; compare_digest avoids timing leaks.
    expected = hashlib.sha1(raw.encode(), usedforsecurity=False).hexdigest()
    return hmac.compare_digest(expected, signature)


def _decrypt_msg(encrypted: str) -> str:
    """Decrypt WeChat Work callback message. Returns plaintext XML."""
    key = base64.b64decode(settings.WECOM_ENCODING_AES_KEY + "=")
    ciphertext = base64.b64decode(encrypted)
    decryptor = Cipher(algorithms.AES(key), modes.CBC(key[:16])).decryptor()
    plaintext = decryptor.update(ciphertext) + decryptor.finalize()
    pad_len = plaintext[-1]
    if pad_len < 1 or pad_len > 32 or plaintext[-pad_len:] != bytes([pad_len]) * pad_len:
        raise ValueError("Invalid callback padding")
    plaintext = plaintext[:-pad_len]
    msg_len = struct.unpack(">I", plaintext[16:20])[0]
    return plaintext[20:20 + msg_len].decode("utf-8")


# ── Callback ──

@router.get("/callback")
async def wecom_callback_verify(
    msg_signature: str = Query(...),
    timestamp: str = Query(...),
    nonce: str = Query(...),
    echostr: str = Query(...),
):
    """WeChat Work callback URL verification (GET)."""
    token = settings.WECOM_TOKEN
    if not token or not settings.WECOM_ENCODING_AES_KEY:
        raise HTTPException(status_code=500, detail="WeCom callback not configured")
    if not _verify_signature(token, timestamp, nonce, echostr, msg_signature):
        raise HTTPException(status_code=403, detail="Signature mismatch")
    try:
        plaintext = _decrypt_msg(echostr)
        return PlainTextResponse(content=plaintext, status_code=200)
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Decrypt failed: {e}")


@router.post("/callback")
async def wecom_callback_event(request: Request):
    """WeChat Work callback event receiver (POST). Handles menu clicks and text messages."""
    token = settings.WECOM_TOKEN
    if not token or not settings.WECOM_ENCODING_AES_KEY:
        return PlainTextResponse(content="success")

    # Parse XML envelope
    try:
        root = ET.fromstring(await request.body())
        encrypt = root.findtext("Encrypt", "")
    except ET.ParseError:
        return PlainTextResponse(content="success")

    # Verify signature
    msg_signature = request.query_params.get("msg_signature", "")
    timestamp = request.query_params.get("timestamp", "")
    nonce = request.query_params.get("nonce", "")
    if not _verify_signature(token, timestamp, nonce, encrypt, msg_signature):
        return PlainTextResponse(content="success")

    # Decrypt inner XML
    try:
        xml_str = _decrypt_msg(encrypt)
        msg = ET.fromstring(xml_str)
    except Exception:
        return PlainTextResponse(content="success")

    msg_type = msg.findtext("MsgType", "")
    from_user = msg.findtext("FromUserName", "")
    content = (msg.findtext("Content", "") or "").strip()
    event = (msg.findtext("Event", "") or "").strip()
    event_key = (msg.findtext("EventKey", "") or "").strip()

    import logging
    _log = logging.getLogger("wecom_callback")
    _log.info("wecom callback received: type=%s event=%s key=%s", msg_type, event, event_key)

    # Triggers for binding
    if (msg_type == "text" and content == "绑定") or (msg_type == "event" and event == "click" and event_key == "BIND_ACCOUNT"):
        _log.info("wecom bind triggered")
        await _handle_bind_request(from_user)
        return PlainTextResponse(content="success")

    _log.info("wecom message ignored (no matching action)")
    return PlainTextResponse(content="success")


async def _handle_bind_request(wecom_userid: str):
    """Generate bind token, store it, and push a bind link to the user."""
    if not wecom_userid:
        return

    async with async_session() as db:
        bind_token = await store_bind_token(db, wecom_userid)

    content = (
        f"【账号绑定】\n\n"
        f"请在10分钟内点击以下链接完成账号绑定：\n"
        f"https://qj.dhdx.fun/wecom-bind?token={bind_token}\n\n"
        f"绑定后可使用企微一键登录，并接收填报提醒通知。"
    )

    await wecom_client.send_text_message([wecom_userid], content)


# ── Bind confirmation (JWT-protected) ──

@router.post("/bind-confirm")
async def bind_confirm(
    data: BindConfirmRequest,
    current_user: dict = Depends(get_current_user),
    db: AsyncSession = Depends(get_db),
):
    """User confirms binding after clicking the link in WeChat Work message."""
    if not data.token:
        raise HTTPException(status_code=400, detail="token 不能为空")

    wecom_userid = await consume_bind_token(db, data.token)
    if not wecom_userid:
        raise HTTPException(status_code=400, detail="绑定链接已过期或无效，请重新在企微发送「绑定」")

    # Check if this wecom_userid is already bound to another user
    result = await db.execute(
        select(User).where(User.wecom_userid == wecom_userid)
    )
    existing = result.scalar_one_or_none()
    if existing and str(existing.id) != current_user["user_id"]:
        raise HTTPException(status_code=400, detail=f"该企业微信已绑定账号「{existing.name}」")

    # Bind
    result = await db.execute(
        select(User).where(User.id == uuid.UUID(current_user["user_id"]))
    )
    user = result.scalar_one()
    user.wecom_userid = wecom_userid
    try:
        await db.commit()
    except IntegrityError as exc:
        await db.rollback()
        raise HTTPException(status_code=409, detail="该企业微信已绑定其他账号") from exc

    return {"code": 200, "msg": "绑定成功", "wecom_userid": wecom_userid}


# ── Manual actions (director only) ──

@router.post("/remind")
async def send_reminder(
    data: RemindRequest,
    current_user: dict = Depends(require_director),
    db: AsyncSession = Depends(get_db),
):
    """Director manually sends reminder to specific managers."""
    result = await db.execute(
        select(User.wecom_userid).where(User.id.in_([uuid.UUID(uid) for uid in data.user_ids]))
    )
    wecom_ids = [r[0] for r in result.all() if r[0]]

    content = data.message or "📋 请及时完成今日拜访记录填报。"
    success = await wecom_client.send_template_card(
        user_ids=wecom_ids,
        title="📋 填报提醒",
        desc=content,
        url="https://qj.dhdx.fun/m",
        btn_text="去填报",
    )
    if not success:
        success = await wecom_client.send_text_message(wecom_ids, content)

    return {"success": success, "sent_to": len(wecom_ids), "rich": True}


@router.post("/announcement")
async def send_announcement(
    data: AnnouncementRequest,
    current_user: dict = Depends(require_director),
    db: AsyncSession = Depends(get_db),
):
    """Director sends an announcement to all team members (markdown)."""
    result = await db.execute(select(User.wecom_userid).where(User.wecom_userid.isnot(None)))
    wecom_ids = [r[0] for r in result.all()]

    markdown = f"## 📢 支局长公告\n\n{data.content}"
    success = await wecom_client.send_markdown_message(markdown)
    if not success:
        success = await wecom_client.send_text_message(wecom_ids, f"📢 支局长公告\n\n{data.content}")

    return {"success": success, "sent_to": len(wecom_ids), "rich": True}


@router.post("/test-message")
async def test_message(
    wecom_userid: str,
    current_user: dict = Depends(require_director),
):
    """Send a test message to a specific wecom_userid for debugging."""
    success = await wecom_client.send_text_message(
        [wecom_userid],
        "🧪 企迹周报系统 — 测试消息\n\n如果您收到此消息，说明企微消息推送配置成功！"
    )
    return {"success": success, "sent_to": wecom_userid}


@router.post("/setup-menu")
async def setup_menu(current_user: dict = Depends(require_director)):
    """Create the standard app menu (开始填报 view + 绑定账号 click)."""
    buttons = [
        {"type": "view", "name": "开始填报", "url": "https://qj.dhdx.fun/m"},
        {"type": "click", "name": "绑定账号", "key": "BIND_ACCOUNT"},
    ]
    success = await wecom_client.create_menu(buttons)
    return {"success": success, "message": "菜单已部署 (view + click)" if success else "菜单创建失败"}


@router.get("/menu")
async def get_menu(current_user: dict = Depends(require_director)):
    """Get current app menu configuration."""
    menu = await wecom_client.get_menu()
    return menu or {"message": "No menu configured"}


@router.post("/trigger-daily-check")
async def trigger_daily_check(
    current_user: dict = Depends(require_director),
    db: AsyncSession = Depends(get_db),
):
    """Manually trigger the daily reporting check (for testing or manual use)."""
    result = await check_daily_reporting(db)
    return result