douyin-chat-export/extract.py at main · TeamBreakerr/douyin-chat-export · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
#!/usr/bin/env python3
"""Main entry point: extract Douyin chat data via web version."""
import asyncio
import sys
import os

sys.path.insert(0, os.path.dirname(__file__))


def _parse_args():
    """Parse CLI arguments."""
    args = {
        "mode": "extract",
        "name_filter": None,
        "incremental": "--incremental" in sys.argv,
        "download_images": "--download-images" in sys.argv,
        "output_format": "jsonl",
        "output_path": None,
    }

    if "--discover" in sys.argv:
        args["mode"] = "discover"
    elif "--list-conversations" in sys.argv:
        args["mode"] = "list_conversations"
    elif "--export" in sys.argv:
        args["mode"] = "export"

    for i, arg in enumerate(sys.argv[1:], 1):
        if arg == "--filter" and i < len(sys.argv) - 1:
            args["name_filter"] = sys.argv[i + 1]
        elif arg == "--format" and i < len(sys.argv) - 1:
            args["output_format"] = sys.argv[i + 1]
        elif arg == "--output" and i < len(sys.argv) - 1:
            args["output_path"] = sys.argv[i + 1]

    return args


def run_export(args):
    """Export chat data to ChatLab format (no browser needed)."""
    from extractor.exporter import ChatLabExporter

    fmt = args["output_format"]
    ext = ".json" if fmt == "json" else ".jsonl"
    output_path = args["output_path"] or os.path.join("data", f"export{ext}")

    exporter = ChatLabExporter(
        conv_name=args["name_filter"],
        output_format=fmt,
    )
    exporter.export(output_path)


async def run():
    args = _parse_args()

    # Export mode: no browser needed
    if args["mode"] == "export":
        run_export(args)
        return 0

    from extractor.web_scraper import WebChatScraper

    scraper = WebChatScraper(
        discovery_mode=(args["mode"] == "discover"),
        name_filter=args["name_filter"],
        incremental=args["incremental"],
        download_images=args["download_images"],
    )

    try:
        await scraper.launch()
        logged_in = await scraper.wait_for_login()
        if not logged_in:
            print("[-] 未能登录，退出")
            return 2  # non-zero exit so the panel surfaces this as a failure

        if args["mode"] == "discover":
            duration = 60
            for arg in sys.argv[1:]:
                if arg.isdigit():
                    duration = int(arg)
            await scraper.run_discovery(duration=duration)
        elif args["mode"] == "list_conversations":
            convs = await scraper.list_conversations()
            out_path = os.path.join(
                os.path.dirname(__file__), "data", "conversations_list.json"
            )
            os.makedirs(os.path.dirname(out_path), exist_ok=True)
            import json as _json
            import time as _time
            payload = {
                "discovered_at": int(_time.time()),
                "items": [
                    {
                        "nickname": c.get("nickname", ""),
                        "name": c.get("name", ""),
                        "time": c.get("time", ""),
                        "preview": c.get("preview", ""),
                    }
                    for c in convs
                ],
            }
            with open(out_path, "w", encoding="utf-8") as f:
                _json.dump(payload, f, ensure_ascii=False, indent=2)
            print(f"[+] 会话列表已写入 {out_path}")
        else:
            await scraper.extract_all()

        return 0

    except KeyboardInterrupt:
        print("\n[*] 用户中断")
        return 130
    except Exception as e:
        print(f"\n[-] 错误: {e}")
        import traceback
        traceback.print_exc()
        return 1
    finally:
        await scraper.close()


if __name__ == "__main__":
    sys.exit(asyncio.run(run()) or 0)