dc-docs/scripts/extract_group_members.py

"""从钉钉群拉取全部成员并查询工号/邮箱/手机等详细信息，输出 CSV"""

import csv
import json
import os
import subprocess
import sys
import time

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__))))
from config import DWS

GROUP_ID = "cidrLgqKZUgbJ9PdVHgQQIS9g=="
OUTPUT_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data")
OUTPUT_CSV = os.path.join(OUTPUT_DIR, "group_members_detail.csv")
OUTPUT_JSON = os.path.join(OUTPUT_DIR, "group_members_all.json")


def run_dws(args, timeout=30):
    cmd = [DWS] + args + ["-f", "json"]
    result = subprocess.run(cmd, capture_output=True, timeout=timeout)
    out = result.stdout.decode("utf-8", errors="replace")
    for i, ch in enumerate(out):
        if ch == "{":
            return json.loads(out[i:])
    return None


def fetch_all_members():
    all_members = []
    cursor = "0"
    page = 0
    while True:
        page += 1
        args = ["chat", "group", "members", "list", "--id", GROUP_ID, "--cursor", cursor]
        data = run_dws(args, timeout=60)
        if not data or not data.get("success"):
            print(f"  page {page} failed: {data}")
            break
        result = data.get("result", {})
        members = result.get("list", [])
        all_members.extend(members)
        print(f"  page {page}: {len(members)} members (total {len(all_members)})")
        if not result.get("hasMore"):
            break
        cursor = str(result.get("nextCursor", ""))
        if not cursor:
            break
        time.sleep(0.3)
    return all_members


def search_user_id(name, open_dingtalk_id):
    """通过姓名搜索获取 userId"""
    args = ["contact", "user", "search", "--query", name]
    data = run_dws(args, timeout=15)
    if data and data.get("success"):
        results = data.get("result", [])
        for r in results:
            if r.get("openDingTalkId") == open_dingtalk_id:
                return r.get("userId", "")
            if r.get("name") == name:
                return r.get("userId", "")
        if results:
            return results[0].get("userId", "")
    return ""


def batch_get_user_details(user_ids):
    """批量获取用户详情"""
    profiles = {}
    batch_size = 20
    for i in range(0, len(user_ids), batch_size):
        batch = [uid for uid in user_ids[i:i+batch_size] if uid]
        if not batch:
            continue
        ids_str = ",".join(batch)
        args = ["contact", "user", "get", "--ids", ids_str]
        data = run_dws(args, timeout=30)
        if data and data.get("success"):
            result = data.get("result", [])
            if isinstance(result, list):
                for u in result:
                    model = u.get("orgEmployeeModel", {})
                    uid = model.get("orgUserId", "")
                    if uid:
                        profiles[uid] = model
        print(f"  batch {i//batch_size+1}/{(len(user_ids)+batch_size-1)//batch_size}: got {len(profiles)} profiles")
        time.sleep(0.3)
    return profiles


def main():
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    print("=" * 50)
    print("Step 1: Fetch all group members")
    print("=" * 50)
    members = fetch_all_members()
    print(f"\nTotal: {len(members)} members")

    with open(OUTPUT_JSON, "w", encoding="utf-8") as f:
        json.dump(members, f, ensure_ascii=False, indent=2)

    print("\n" + "=" * 50)
    print("Step 2: Resolve userId for each member via search")
    print("=" * 50)
    user_id_map = {}  # openDingtalkId -> userId
    for idx, m in enumerate(members, 1):
        name = m.get("memberEmpName", "")
        oid = m.get("openDingtalkId", "")
        if name and oid:
            uid = search_user_id(name, oid)
            if uid:
                user_id_map[oid] = uid
        if idx % 50 == 0:
            print(f"  resolved {len(user_id_map)}/{idx} userIds")
        time.sleep(0.15)
    print(f"  resolved {len(user_id_map)}/{len(members)} userIds total")

    print("\n" + "=" * 50)
    print("Step 3: Batch query user profiles")
    print("=" * 50)
    all_user_ids = list(user_id_map.values())
    profiles = batch_get_user_details(all_user_ids)

    print("\n" + "=" * 50)
    print("Step 4: Generate CSV")
    print("=" * 50)

    fieldnames = ["序号", "姓名", "昵称", "工号", "企业邮箱", "职位", "部门", "userId", "角色", "openDingtalkId"]
    rows = []
    for idx, m in enumerate(members, 1):
        oid = m.get("openDingtalkId", "")
        name = m.get("memberEmpName", "")
        nick = m.get("memberNick", "")
        role = m.get("memberRoleDesc", "")

        uid = user_id_map.get(oid, "")
        p = profiles.get(uid, {})
        job_number = p.get("jobNumber", "")
        email = p.get("orgAuthEmail", "")
        title = p.get("orgTitle", "")
        depts = p.get("depts", [])
        dept_name = "/".join(d.get("deptName", "") for d in depts) if depts else ""

        rows.append({
            "序号": idx, "姓名": name, "昵称": nick, "工号": job_number,
            "企业邮箱": email, "职位": title, "部门": dept_name,
            "userId": uid, "角色": role, "openDingtalkId": oid,
        })

    with open(OUTPUT_CSV, "w", encoding="utf-8-sig", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)

    print(f"\nCSV saved: {OUTPUT_CSV}")
    print(f"Total: {len(rows)} rows")
    has_job = sum(1 for r in rows if r["工号"])
    has_email = sum(1 for r in rows if r["企业邮箱"])
    has_dept = sum(1 for r in rows if r["部门"])
    print(f"有工号: {has_job}, 有邮箱: {has_email}, 有部门: {has_dept}")


if __name__ == "__main__":
    main()