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, HTMLResponse
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)."""
    oauth_url = _build_oauth_url()
    buttons = [
        {"type": "view", "name": "开始填报", "url": oauth_url},
        {"type": "click", "name": "绑定账号", "key": "BIND_ACCOUNT"},
    ]
    success = await wecom_client.create_menu(buttons)
    return {"success": success, "message": "菜单已部署 (含静默登录)" 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


# ── OAuth silent login ──

import urllib.parse
from app.services.auth import build_token_for_user


def _build_oauth_url(redirect_path: str = "/m") -> str:
    """Build the WeChat Work OAuth2 authorize URL (snsapi_base = silent).
    The redirect target is encoded in the 'state' param to avoid query-string issues
    with the OAuth redirect_uri validation."""
    if not settings.WECOM_CORP_ID:
        return f"https://qj.dhdx.fun{redirect_path}"
    params = urllib.parse.urlencode({
        "appid": settings.WECOM_CORP_ID,
        "redirect_uri": "https://qj.dhdx.fun/api/wecom/oauth-callback",
        "response_type": "code",
        "scope": "snsapi_base",
        "agentid": settings.WECOM_AGENT_ID,
        "state": f"r={urllib.parse.quote(redirect_path)}",
    })
    return f"https://open.weixin.qq.com/connect/oauth2/authorize?{params}#wechat_redirect"


@router.get("/oauth-url")
async def get_oauth_url(redirect: str = Query("/m")):
    """Redirect to the WeChat Work OAuth authorize URL (silent login)."""
    from fastapi.responses import RedirectResponse
    return RedirectResponse(url=_build_oauth_url(redirect))


@router.get("/oauth-callback")
async def oauth_callback(
    code: str = Query(...),
    state: str = Query(""),
    db: AsyncSession = Depends(get_db),
):
    """WeChat Work OAuth2 callback — exchange code for userid, find bound user, auto-login."""
    # Parse redirect target from state param: "r=%2Fm"
    redirect = "/m"
    if state.startswith("r="):
        redirect = urllib.parse.unquote(state[2:])

    # Exchange code for wecom userinfo
    userinfo = await wecom_client.get_userinfo_by_code(code)
    if not userinfo:
        raise HTTPException(status_code=400, detail="Failed to exchange code with WeChat Work")

    wecom_userid = userinfo.get("UserId") or userinfo.get("userid")
    if not wecom_userid:
        raise HTTPException(status_code=400, detail="Could not get userid from WeChat Work")

    # Find bound user (case-insensitive — WeChat Work userids may vary in case)
    from sqlalchemy import func as sa_func
    result = await db.execute(
        select(User).where(sa_func.lower(User.wecom_userid) == wecom_userid.lower())
    )
    user = result.scalar_one_or_none()
    if not user:
        # Not bound — generate a one-time bind token and redirect to bind page
        bind_token = await store_bind_token(db, wecom_userid)
        bind_url = f"https://qj.dhdx.fun/wecom-bind?token={bind_token}&wecom_userid={wecom_userid}"
        from fastapi.responses import RedirectResponse
        return RedirectResponse(url=bind_url)

    # Issue JWT and auto-login via HTML page — saves credentials to localStorage, then redirects
    token = build_token_for_user(user)
    import json as _json
    auth_data = _json.dumps({
        "token": token,
        "userId": str(user.id),
        "userName": user.name,
        "userRole": user.role,
        "theme": user.theme or "editorial",
    }, ensure_ascii=False)
    frontend_url = f"https://qj.dhdx.fun{redirect}"
    html = f"""<!DOCTYPE html><html><head><meta charset="utf-8"><title>登录中...</title></head>
<body><script>
var d = {auth_data};
localStorage.setItem('token', d.token);
localStorage.setItem('userId', d.userId);
localStorage.setItem('userName', d.userName);
localStorage.setItem('userRole', d.userRole);
localStorage.setItem('theme', d.theme);
window.location.replace({_json.dumps(frontend_url)});
</script></body></html>"""
    return HTMLResponse(content=html, status_code=200)