lami/lami-base-v1/scripts/qiwei_doc_discovery.py

"""
QiWei：从「会话消息 / 历史同步 / 回调存档」中间接发现在线文档（提取 docid）。

依据 doc.qiweapi.com：
  - 回调 cmd=15000，链接消息 msgType=13，字段 linkUrl（见回调结构说明）
  - 同步历史消息：method=/msg/syncMsg（见 api-344613926）

用法：
  # 1) 离线：解析已保存的回调 JSON
  python scripts/qiwei_doc_discovery.py analyze --file scripts/samples/qiwei_callbacks_sample.json

  # 2) 在线：拉历史消息并扫描（需环境变量，见 doc/会话存档-文档发现-测试指南.md）
  python scripts/qiwei_doc_discovery.py sync --guid YOUR_GUID --msg-seq 0

环境变量（sync 子命令）：
  QIWEI_TOKEN          必填
  QIWEI_BASE_URL       默认 http://manager.qiweapi.com/qiwe
  QIWEI_SYNC_METHOD    默认 /msg/syncMsg
"""
from __future__ import annotations

import argparse
import json
import os
import re
import sys
from dataclasses import asdict, dataclass
from typing import Any, Iterator
from urllib.parse import parse_qs, urlparse

import requests

# 企微在线文档 / 微盘链接常见形态
DOC_URL_PATTERNS = (
    re.compile(r"doc\.weixin\.qq\.com", re.I),
    re.compile(r"doc\.work\.weixin\.qq\.com", re.I),
    re.compile(r"wedoc", re.I),
    re.compile(r"txdoc", re.I),
)

DOCID_QUERY_KEYS = ("docid", "doc_id", "docId")
DOCID_PATH_RE = re.compile(
    r"(?:docid|doc_id)[=/]([A-Za-z0-9_\-]+)", re.I
)
# 如 https://doc.weixin.qq.com/doc/DOC_OTHER_room2
DOCID_PATH_SEGMENT_RE = re.compile(
    r"doc\.weixin\.qq\.com/(?:txdoc/[^/?#]+|doc)/([A-Za-z0-9_\-]+)", re.I
)


@dataclass
class DocRef:
    docid: str | None
    link_url: str
    title: str | None
    msg_type: int | None
    room_id: str | int | None
    guid: str | None
    source: str  # callback | sync | file


def extract_docid_from_url(url: str) -> str | None:
    if not url or not isinstance(url, str):
        return None
    parsed = urlparse(url.strip())
    qs = parse_qs(parsed.query)
    for key in DOCID_QUERY_KEYS:
        if key in qs and qs[key]:
            return qs[key][0]
    m = DOCID_PATH_RE.search(url)
    if m:
        return m.group(1)
    m = DOCID_PATH_SEGMENT_RE.search(url)
    if m:
        return m.group(1)
    # 部分链接 docid 在 fragment
    if parsed.fragment:
        m = DOCID_PATH_RE.search(parsed.fragment)
        if m:
            return m.group(1)
    return None


def is_probable_doc_url(url: str) -> bool:
    if not url:
        return False
    return any(p.search(url) for p in DOC_URL_PATTERNS)


def _walk(obj: Any) -> Iterator[tuple[str, Any]]:
    if isinstance(obj, dict):
        for k, v in obj.items():
            yield k, v
            yield from _walk(v)
    elif isinstance(obj, list):
        for item in obj:
            yield from _walk(item)


def extract_from_link_msgdata(msg_data: dict) -> DocRef | None:
    link = msg_data.get("linkUrl") or msg_data.get("link_url")
    if not link or not is_probable_doc_url(str(link)):
        if link and "doc.weixin" not in str(link).lower():
            return None
        if not link:
            return None
    title = msg_data.get("title")
    if isinstance(title, str) and title.isascii() and len(title) > 40:
        try:
            import base64

            title = base64.b64decode(title).decode("utf-8", errors="replace")
        except Exception:
            pass
    return DocRef(
        docid=extract_docid_from_url(str(link)),
        link_url=str(link),
        title=str(title) if title else None,
        msg_type=13,
        room_id=None,
        guid=None,
        source="msgData",
    )


def extract_doc_refs_from_message_item(item: dict, source: str) -> list[DocRef]:
    refs: list[DocRef] = []
    msg_type = item.get("msgType")
    room_id = item.get("fromRoomId") or item.get("roomId") or item.get("roomid")
    guid = item.get("guid")
    msg_data = item.get("msgData")

    if isinstance(msg_data, dict):
        if msg_type == 13 or "linkUrl" in msg_data or "link_url" in msg_data:
            ref = extract_from_link_msgdata(msg_data)
            if ref:
                ref.msg_type = msg_type
                ref.room_id = room_id
                ref.guid = guid
                ref.source = source
                refs.append(ref)

    # 兜底：整棵 JSON 里搜 doc.weixin 链接
    for key, val in _walk(item):
        if key.lower() in ("linkurl", "link_url", "url", "link") and isinstance(val, str):
            if is_probable_doc_url(val) or "doc.weixin" in val.lower():
                refs.append(
                    DocRef(
                        docid=extract_docid_from_url(val),
                        link_url=val,
                        title=None,
                        msg_type=msg_type if isinstance(msg_type, int) else None,
                        room_id=room_id,
                        guid=guid,
                        source=f"{source}:walk",
                    )
                )
    return refs


def extract_doc_refs_from_payload(payload: Any, source: str = "callback") -> list[DocRef]:
    refs: list[DocRef] = []
    if isinstance(payload, dict):
        data = payload.get("data")
        if isinstance(data, list):
            for item in data:
                if isinstance(item, dict):
                    refs.extend(extract_doc_refs_from_message_item(item, source))
        else:
            refs.extend(extract_doc_refs_from_message_item(payload, source))
    elif isinstance(payload, list):
        for item in payload:
            refs.extend(extract_doc_refs_from_payload(item, source))
    # 去重（按 link_url）
    seen: set[str] = set()
    unique: list[DocRef] = []
    for r in refs:
        if r.link_url not in seen:
            seen.add(r.link_url)
            unique.append(r)
    return unique


def qiwei_do_api(
    token: str,
    method: str,
    params: dict,
    base_url: str | None = None,
) -> dict:
    base = (base_url or os.environ.get("QIWEI_BASE_URL") or "http://manager.qiweapi.com/qiwe").rstrip("/")
    url = f"{base}/api/qw/doApi"
    headers = {
        "Content-Type": "application/json",
        "X-QIWEI-TOKEN": token,
    }
    body = {"method": method, "params": params}
    resp = requests.post(url, headers=headers, json=body, timeout=60)
    resp.raise_for_status()
    return resp.json()


def sync_messages_and_scan(
    token: str,
    guid: str,
    msg_seq: int = 0,
    *,
    base_url: str | None = None,
    sync_method: str | None = None,
    extra_params: dict | None = None,
) -> tuple[dict, list[DocRef]]:
    method = sync_method or os.environ.get("QIWEI_SYNC_METHOD", "/msg/syncMsg")
    params: dict[str, Any] = {"guid": guid, "msgSeq": msg_seq}
    if extra_params:
        params.update(extra_params)
    raw = qiwei_do_api(token, method, params, base_url)
    refs: list[DocRef] = []
    # 响应里可能是 data.syncMsgList 或 data 为列表
    data = raw.get("data")
    if isinstance(data, dict):
        lst = data.get("syncMsgList") or data.get("list") or []
        if isinstance(lst, list):
            for item in lst:
                if isinstance(item, dict):
                    refs.extend(extract_doc_refs_from_message_item(item, "sync"))
    elif isinstance(data, list):
        for item in data:
            if isinstance(item, dict):
                refs.extend(extract_doc_refs_from_message_item(item, "sync"))
    refs.extend(extract_doc_refs_from_payload(raw, "sync:raw"))
    return raw, refs


def cmd_analyze(args: argparse.Namespace) -> int:
    path = args.file
    with open(path, encoding="utf-8") as f:
        payload = json.load(f)
    refs = extract_doc_refs_from_payload(payload, "file")
    print(f"文件: {path}")
    print(f"发现文档相关消息: {len(refs)} 条\n")
    for i, r in enumerate(refs, 1):
        print(f"--- [{i}] ---")
        print(json.dumps(asdict(r), ensure_ascii=False, indent=2))
    if not refs:
        print("未发现含 doc.weixin / docid 的链接。请确认样本中含 msgType=13 且 linkUrl 指向在线文档。")
        return 1
    missing = [r for r in refs if not r.docid]
    if missing:
        print(f"\n警告: {len(missing)} 条链接未能解析出 docid，需人工从 link_url 查看。")
    return 0


def cmd_sync(args: argparse.Namespace) -> int:
    token = os.environ.get("QIWEI_TOKEN")
    if not token:
        print("请设置环境变量 QIWEI_TOKEN（控制台租户令牌）", file=sys.stderr)
        return 2
    guid = args.guid or os.environ.get("QIWEI_GUID")
    if not guid:
        print("请传入 --guid 或设置 QIWEI_GUID", file=sys.stderr)
        return 2
    extra = {}
    if args.room_id:
        extra["fromRoomId"] = args.room_id
    raw, refs = sync_messages_and_scan(
        token,
        guid,
        args.msg_seq,
        extra_params=extra or None,
    )
    out_raw = args.out_raw or "output/qiwei_sync_last.json"
    os.makedirs(os.path.dirname(out_raw) or ".", exist_ok=True)
    with open(out_raw, "w", encoding="utf-8") as f:
        json.dump(raw, f, ensure_ascii=False, indent=2)
    print(f"原始响应已保存: {out_raw}")
    print(f"扫描到文档相关消息: {len(refs)} 条\n")
    for r in refs:
        print(json.dumps(asdict(r), ensure_ascii=False))
    return 0 if refs else 1


def main() -> int:
    parser = argparse.ArgumentParser(description="QiWei 会话消息中间接发现在线文档（docid）")
    sub = parser.add_subparsers(dest="command", required=True)

    p_analyze = sub.add_parser("analyze", help="解析本地回调/存档 JSON")
    p_analyze.add_argument("--file", "-f", required=True, help="JSON 文件路径")
    p_analyze.set_defaults(func=cmd_analyze)

    p_sync = sub.add_parser("sync", help="调用 /msg/syncMsg 拉历史并扫描")
    p_sync.add_argument("--guid", help="账号 guid")
    p_sync.add_argument("--msg-seq", type=int, default=0, help="起始 msgSeq，首次可传 0")
    p_sync.add_argument("--room-id", help="可选：仅关注某群 fromRoomId")
    p_sync.add_argument("--out-raw", help="保存完整 API 响应的路径")
    p_sync.set_defaults(func=cmd_sync)

    args = parser.parse_args()
    return args.func(args)


if __name__ == "__main__":
    raise SystemExit(main())