| 1 | import argparse |
| 2 | import asyncio |
| 3 | import json |
| 4 | import logging |
| 5 | import sys |
| 6 | from pathlib import Path |
| 7 | from typing import Any |
| 8 | |
| 9 | from auth import CookieManager |
| 10 | from cli.progress_display import ProgressDisplay |
| 11 | from config import ConfigLoader |
| 12 | from control import QueueManager, RateLimiter, RetryHandler |
| 13 | from core import DouyinAPIClient, DownloaderFactory, URLParser |
| 14 | from storage import Database, FileManager |
| 15 | from utils.logger import set_console_log_level, setup_logger |
| 16 | from utils.notifier import build_notifier |
| 17 | from utils.validators import is_short_url, normalize_short_url |
| 18 | |
| 19 | logger = setup_logger("CLI") |
| 20 | display = ProgressDisplay() |
| 21 | |
| 22 | |
| 23 | def _as_bool(value: Any, default: bool = True) -> bool: |
| 24 | if value is None: |
| 25 | return default |
| 26 | if isinstance(value, bool): |
| 27 | return value |
| 28 | if isinstance(value, str): |
| 29 | return value.strip().lower() in {"1", "true", "yes", "on"} |
| 30 | return bool(value) |
| 31 | |
| 32 | |
| 33 | async def download_url( |
| 34 | url: str, |
| 35 | config: ConfigLoader, |
| 36 | cookie_manager: CookieManager, |
| 37 | database: Database = None, |
| 38 | progress_reporter: ProgressDisplay = None, |
| 39 | ): |
| 40 | if progress_reporter: |
| 41 | progress_reporter.advance_step("初始化", "创建下载组件") |
| 42 | file_manager = FileManager(config.get("path")) |
| 43 | rate_limiter = RateLimiter(max_per_second=float(config.get("rate_limit", 2) or 2)) |
| 44 | retry_handler = RetryHandler(max_retries=config.get("retry_times", 3)) |
| 45 | queue_manager = QueueManager(max_workers=int(config.get("thread", 5) or 5)) |
| 46 | |
| 47 | original_url = url |
| 48 | |
| 49 | async with DouyinAPIClient( |
| 50 | cookie_manager.get_cookies(), |
| 51 | proxy=config.get("proxy"), |
| 52 | ) as api_client: |
| 53 | if progress_reporter: |
| 54 | progress_reporter.advance_step("解析链接", "检查短链并解析 URL") |
| 55 | # 支持多种短链变体：v.douyin.com / v.iesdouyin.com / 无 scheme 的裸链接 |
| 56 | if is_short_url(url): |
| 57 | resolved_url = await api_client.resolve_short_url(normalize_short_url(url)) |
| 58 | if resolved_url: |
| 59 | url = resolved_url |
| 60 | else: |
| 61 | if progress_reporter: |
| 62 | progress_reporter.update_step("解析链接", "短链解析失败") |
| 63 | display.print_error(f"Failed to resolve short URL: {url}") |
| 64 | return None |
| 65 | |
| 66 | parsed = URLParser.parse(url) |
| 67 | if not parsed: |
| 68 | if progress_reporter: |
| 69 | progress_reporter.update_step("解析链接", "URL 解析失败") |
| 70 | display.print_error(f"Failed to parse URL: {url}") |
| 71 | return None |
| 72 | |
| 73 | if not progress_reporter: |
| 74 | display.print_info(f"URL type: {parsed['type']}") |
| 75 | if progress_reporter: |
| 76 | progress_reporter.advance_step("创建下载器", f"URL 类型: {parsed['type']}") |
| 77 | |
| 78 | downloader = DownloaderFactory.create( |
| 79 | parsed["type"], |
| 80 | config, |
| 81 | api_client, |
| 82 | file_manager, |
| 83 | cookie_manager, |
| 84 | database, |
| 85 | rate_limiter, |
| 86 | retry_handler, |
| 87 | queue_manager, |
| 88 | progress_reporter=progress_reporter, |
| 89 | ) |
| 90 | |
| 91 | if not downloader: |
| 92 | if progress_reporter: |
| 93 | progress_reporter.update_step("创建下载器", "未找到匹配下载器") |
| 94 | display.print_error(f"No downloader found for type: {parsed['type']}") |
| 95 | return None |
| 96 | |
| 97 | if progress_reporter: |
| 98 | progress_reporter.advance_step("执行下载", "开始拉取与下载资源") |
| 99 | try: |
| 100 | result = await downloader.download(parsed) |
| 101 | except Exception as exc: |
| 102 | # Surface fatal downloader errors (e.g. user_info fetch failed |
| 103 | # because cookies are invalid) as a per-URL failure instead of |
| 104 | # crashing the whole batch. Keeps multi-URL CLI runs robust while |
| 105 | # still telling the user why the URL was skipped. |
| 106 | if progress_reporter: |
| 107 | progress_reporter.update_step("执行下载", f"失败：{exc}") |
| 108 | display.print_error(f"Download failed for {url}: {exc}") |
| 109 | return None |
| 110 | |
| 111 | if progress_reporter: |
| 112 | progress_reporter.advance_step( |
| 113 | "记录历史", |
| 114 | "写入数据库历史" if (result and database) else "数据库未启用，跳过", |
| 115 | ) |
| 116 | if result and database: |
| 117 | safe_config = { |
| 118 | k: v |
| 119 | for k, v in config.config.items() |
| 120 | if k not in ("cookies", "cookie", "transcript") |
| 121 | } |
| 122 | await database.add_history( |
| 123 | { |
| 124 | "url": original_url, |
| 125 | "url_type": parsed["type"], |
| 126 | "total_count": result.total, |
| 127 | "success_count": result.success, |
| 128 | "config": json.dumps(safe_config, ensure_ascii=False), |
| 129 | } |
| 130 | ) |
| 131 | |
| 132 | if progress_reporter: |
| 133 | if result: |
| 134 | progress_reporter.advance_step( |
| 135 | "收尾", |
| 136 | f"成功 {result.success} / 失败 {result.failed} / 跳过 {result.skipped}", |
| 137 | ) |
| 138 | else: |
| 139 | progress_reporter.advance_step("收尾", "无可统计结果") |
| 140 | |
| 141 | return result |
| 142 | |
| 143 | |
| 144 | async def main_async(args): |
| 145 | if not args.serve: |
| 146 | display.show_banner() |
| 147 | |
| 148 | if args.config: |
| 149 | config_path = args.config |
| 150 | else: |
| 151 | config_path = "config.yml" |
| 152 | |
| 153 | # 若 config 不存在且使用了 --hot-board / --search / --serve 等独立子命令， |
| 154 | # 允许以默认配置运行（只要命令行提供了 --path）。 |
| 155 | if not Path(config_path).exists(): |
| 156 | if not (args.hot_board is not None or args.search or args.serve): |
| 157 | display.print_error(f"Config file not found: {config_path}") |
| 158 | return |
| 159 | # For ``--serve`` we still pass the (yet-missing) path so later |
| 160 | # ``config.save()`` calls from the REST settings endpoint create |
| 161 | # the file in the right place (e.g. Electron's userData dir). |
| 162 | # Other subcommands keep the historical behaviour of in-memory |
| 163 | # defaults. |
| 164 | if args.serve and args.config: |
| 165 | config = ConfigLoader(config_path) |
| 166 | else: |
| 167 | config = ConfigLoader(None) |
| 168 | else: |
| 169 | config = ConfigLoader(config_path) |
| 170 | |
| 171 | if args.path: |
| 172 | config.update(path=args.path) |
| 173 | |
| 174 | # 独立子命令：热榜 / 搜索 / 服务 |
| 175 | if args.hot_board is not None or args.search: |
| 176 | await _run_discovery_subcommand(args, config) |
| 177 | return |
| 178 | if args.serve: |
| 179 | await _run_serve_subcommand(args, config) |
| 180 | return |
| 181 | |
| 182 | if args.url: |
| 183 | urls = args.url if isinstance(args.url, list) else [args.url] |
| 184 | for url in urls: |
| 185 | if url not in config.get("link", []): |
| 186 | config.update(link=config.get("link", []) + [url]) |
| 187 | |
| 188 | if args.thread: |
| 189 | config.update(thread=args.thread) |
| 190 | |
| 191 | if not config.validate(): |
| 192 | display.print_error("Invalid configuration: missing required fields") |
| 193 | return |
| 194 | |
| 195 | cookies = config.get_cookies() |
| 196 | cookie_manager = CookieManager() |
| 197 | cookie_manager.set_cookies(cookies) |
| 198 | |
| 199 | if not cookie_manager.validate_cookies(): |
| 200 | display.print_warning("Cookies may be invalid or incomplete") |
| 201 | |
| 202 | database = None |
| 203 | if config.get("database"): |
| 204 | db_path = config.get("database_path", "dy_downloader.db") or "dy_downloader.db" |
| 205 | database = Database(db_path=str(db_path)) |
| 206 | await database.initialize() |
| 207 | display.print_success("Database initialized") |
| 208 | |
| 209 | urls = config.get_links() |
| 210 | display.print_info(f"Found {len(urls)} URL(s) to process") |
| 211 | |
| 212 | all_results = [] |
| 213 | progress_config = config.get("progress", {}) or {} |
| 214 | quiet_by_config = _as_bool(progress_config.get("quiet_logs", True), default=True) |
| 215 | quiet_progress_logs = quiet_by_config and not (args.verbose or args.show_warnings) |
| 216 | if quiet_progress_logs: |
| 217 | # Progress 运行期间若有大量错误日志会触发 rich 反复重绘，导致屏幕出现重复块。 |
| 218 | # 默认静默控制台日志，下载完成后再恢复。 |
| 219 | set_console_log_level(logging.CRITICAL) |
| 220 | |
| 221 | display.start_download_session(len(urls)) |
| 222 | try: |
| 223 | for i, url in enumerate(urls, 1): |
| 224 | display.start_url(i, len(urls), url) |
| 225 | |
| 226 | result = await download_url( |
| 227 | url, |
| 228 | config, |
| 229 | cookie_manager, |
| 230 | database, |
| 231 | progress_reporter=display, |
| 232 | ) |
| 233 | if result: |
| 234 | all_results.append(result) |
| 235 | display.complete_url(result) |
| 236 | else: |
| 237 | display.fail_url("下载失败或链接无效") |
| 238 | finally: |
| 239 | display.stop_download_session() |
| 240 | if database is not None: |
| 241 | await database.close() |
| 242 | if quiet_progress_logs: |
| 243 | set_console_log_level(logging.ERROR) |
| 244 | |
| 245 | if all_results: |
| 246 | from core.downloader_base import DownloadResult |
| 247 | |
| 248 | total_result = DownloadResult() |
| 249 | for r in all_results: |
| 250 | total_result.total += r.total |
| 251 | total_result.success += r.success |
| 252 | total_result.failed += r.failed |
| 253 | total_result.skipped += r.skipped |
| 254 | |
| 255 | display.print_success("\n=== Overall Summary ===") |
| 256 | display.show_result(total_result) |
| 257 | |
| 258 | await _dispatch_notifications(config, total_result, len(urls)) |
| 259 | else: |
| 260 | # 所有链接都失败时，也发通知（若启用） |
| 261 | await _dispatch_notifications(config, None, len(urls)) |
| 262 | |
| 263 | |
| 264 | async def _run_discovery_subcommand(args, config: ConfigLoader) -> None: |
| 265 | """处理 --hot-board 与 --search 子命令。""" |
| 266 | from core.discovery import dump_hot_board, search_and_dump |
| 267 | |
| 268 | cookies = config.get_cookies() |
| 269 | cookie_manager = CookieManager() |
| 270 | cookie_manager.set_cookies(cookies) |
| 271 | |
| 272 | base_path = Path(config.get("path") or "./Downloaded/") |
| 273 | |
| 274 | async with DouyinAPIClient(cookie_manager.get_cookies()) as api_client: |
| 275 | if args.hot_board is not None: |
| 276 | display.print_info("拉取抖音热搜榜...") |
| 277 | result = await dump_hot_board(api_client, base_path, limit=int(args.hot_board or 0)) |
| 278 | display.print_success(f"热榜已保存：{result['count']} 条 -> {result['path']}") |
| 279 | if args.search: |
| 280 | display.print_info(f"搜索关键词：{args.search}") |
| 281 | result = await search_and_dump( |
| 282 | api_client, |
| 283 | args.search, |
| 284 | base_path, |
| 285 | max_items=int(args.search_max or 50), |
| 286 | ) |
| 287 | display.print_success(f"搜索结果已保存：{result['count']} 条 -> {result['path']}") |
| 288 | |
| 289 | |
| 290 | async def _run_serve_subcommand(args, config: ConfigLoader) -> None: |
| 291 | """启动 REST API 服务模式（fastapi + uvicorn 为可选依赖）。""" |
| 292 | try: |
| 293 | from server.app import run_server |
| 294 | except ImportError as exc: |
| 295 | display.print_error( |
| 296 | f"REST 服务模式需要安装可选依赖 fastapi + uvicorn：" |
| 297 | f"\n pip install fastapi uvicorn\n原始错误：{exc}" |
| 298 | ) |
| 299 | return |
| 300 | |
| 301 | display.print_info(f"启动 REST 服务：http://{args.serve_host}:{args.serve_port}") |
| 302 | await run_server(config, host=args.serve_host, port=args.serve_port) |
| 303 | |
| 304 | |
| 305 | async def _dispatch_notifications(config: ConfigLoader, total_result: Any, url_count: int) -> None: |
| 306 | notifier = build_notifier(config) |
| 307 | if not notifier.enabled: |
| 308 | return |
| 309 | |
| 310 | if total_result is None: |
| 311 | title = "抖音下载器：全部失败" |
| 312 | body = f"共处理 {url_count} 个链接，无成功结果" |
| 313 | level = "failure" |
| 314 | else: |
| 315 | fail_or_partial = total_result.failed > 0 or total_result.success == 0 |
| 316 | level = "failure" if fail_or_partial else "success" |
| 317 | title = "抖音下载完成" if level == "success" else "抖音下载部分失败" |
| 318 | body = ( |
| 319 | f"链接 {url_count} / 总作品 {total_result.total} / " |
| 320 | f"成功 {total_result.success} / 失败 {total_result.failed} / " |
| 321 | f"跳过 {total_result.skipped}" |
| 322 | ) |
| 323 | |
| 324 | try: |
| 325 | summary = await notifier.send(title=title, body=body, level=level) |
| 326 | if summary: |
| 327 | succ = sum(1 for ok in summary.values() if ok) |
| 328 | logger.info( |
| 329 | "Notification dispatched to %d provider(s), %d ok", |
| 330 | len(summary), |
| 331 | succ, |
| 332 | ) |
| 333 | except Exception as exc: # 通知失败不应影响主流程 |
| 334 | logger.warning("Notification dispatch error: %s", exc) |
| 335 | |
| 336 | |
| 337 | def main(): |
| 338 | parser = argparse.ArgumentParser(description="Douyin Downloader - 抖音批量下载工具") |
| 339 | parser.add_argument("-u", "--url", action="append", help="Download URL(s)") |
| 340 | parser.add_argument("-c", "--config", help="Config file path (default: config.yml)") |
| 341 | parser.add_argument("-p", "--path", help="Save path") |
| 342 | parser.add_argument("-t", "--thread", type=int, help="Thread count") |
| 343 | parser.add_argument("--show-warnings", action="store_true", help="Show warning logs in console") |
| 344 | parser.add_argument("-v", "--verbose", action="store_true", help="Enable verbose console logs") |
| 345 | parser.add_argument( |
| 346 | "--hot-board", |
| 347 | type=int, |
| 348 | nargs="?", |
| 349 | const=0, |
| 350 | default=None, |
| 351 | metavar="N", |
| 352 | help="拉取抖音热搜榜并导出 JSONL，可选上限 N（默认全部）", |
| 353 | ) |
| 354 | parser.add_argument( |
| 355 | "--search", |
| 356 | type=str, |
| 357 | default=None, |
| 358 | metavar="KEYWORD", |
| 359 | help="按关键词搜索作品并导出 JSONL", |
| 360 | ) |
| 361 | parser.add_argument( |
| 362 | "--search-max", |
| 363 | type=int, |
| 364 | default=50, |
| 365 | help="--search 场景下最多拉取条数（默认 50）", |
| 366 | ) |
| 367 | parser.add_argument( |
| 368 | "--serve", |
| 369 | action="store_true", |
| 370 | help="以 REST API 服务模式运行（需要安装 fastapi + uvicorn）", |
| 371 | ) |
| 372 | parser.add_argument("--serve-host", type=str, default="127.0.0.1", help="REST 服务监听地址") |
| 373 | parser.add_argument("--serve-port", type=int, default=8000, help="REST 服务监听端口") |
| 374 | try: |
| 375 | from __init__ import __version__ |
| 376 | except ImportError: |
| 377 | __version__ = "2.0.0" |
| 378 | parser.add_argument("--version", action="version", version=__version__) |
| 379 | |
| 380 | args = parser.parse_args() |
| 381 | |
| 382 | if args.verbose: |
| 383 | set_console_log_level(logging.INFO) |
| 384 | elif args.show_warnings: |
| 385 | set_console_log_level(logging.WARNING) |
| 386 | else: |
| 387 | set_console_log_level(logging.ERROR) |
| 388 | |
| 389 | try: |
| 390 | asyncio.run(main_async(args)) |
| 391 | except KeyboardInterrupt: |
| 392 | display.print_warning("\nDownload interrupted by user") |
| 393 | sys.exit(0) |
| 394 | except Exception as e: |
| 395 | display.print_error(f"Fatal error: {e}") |
| 396 | logger.exception("Fatal error occurred") |
| 397 | sys.exit(1) |
| 398 | |
| 399 | |
| 400 | if __name__ == "__main__": |
| 401 | main() |
| 402 |