| 1 | from __future__ import annotations |
| 2 | |
| 3 | from typing import Any, Dict, List, Optional, Set |
| 4 | |
| 5 | from core.downloader_base import BaseDownloader, DownloadResult |
| 6 | from core.user_mode_registry import UserModeRegistry |
| 7 | from utils.logger import setup_logger |
| 8 | |
| 9 | logger = setup_logger("UserDownloader") |
| 10 | |
| 11 | |
| 12 | class UserDownloader(BaseDownloader): |
| 13 | SELF_COLLECT_MODES = {"collect", "collectmix"} |
| 14 | |
| 15 | def __init__(self, *args, **kwargs): |
| 16 | super().__init__(*args, **kwargs) |
| 17 | self.mode_registry = UserModeRegistry() |
| 18 | self._mode_strategy_cache: Dict[str, Any] = {} |
| 19 | |
| 20 | async def download(self, parsed_url: Dict[str, Any]) -> DownloadResult: |
| 21 | result = DownloadResult() |
| 22 | |
| 23 | sec_uid = parsed_url.get("sec_uid") |
| 24 | if not sec_uid: |
| 25 | # URL parser already validates this; treat as fatal instead of |
| 26 | # a silent empty result so the UI surfaces a real error rather |
| 27 | # than "已完成 0 项". |
| 28 | raise RuntimeError("无法从链接中解析出用户 ID，请确认链接是否完整") |
| 29 | |
| 30 | modes_config = self.config.get("mode", ["post"]) |
| 31 | if isinstance(modes_config, str): |
| 32 | modes = [modes_config] |
| 33 | elif isinstance(modes_config, list): |
| 34 | modes = [str(mode).strip() for mode in modes_config if str(mode).strip()] |
| 35 | else: |
| 36 | modes = ["post"] |
| 37 | |
| 38 | if not self._validate_mode_scope(sec_uid, modes): |
| 39 | return result |
| 40 | |
| 41 | user_info = await self._resolve_user_info(sec_uid, modes) |
| 42 | if not user_info: |
| 43 | logger.error("Failed to get user info: %s", sec_uid) |
| 44 | # Raising here instead of returning an empty result means the |
| 45 | # job ends in `failed` state with a clear message. Returning |
| 46 | # {total:0,success:0,failed:0} made JobManager mark it as |
| 47 | # `success`, which rendered as "已完成 0 项" — a silent failure |
| 48 | # that's indistinguishable from "nothing happened" in the UI. |
| 49 | raise RuntimeError("获取用户信息失败，请检查 Cookie 是否有效或重新登录抖音") |
| 50 | |
| 51 | # Cache author metadata on the hosting job so retry doesn't have |
| 52 | # to re-fetch user_info, and so JobRow can display the nickname. |
| 53 | self._progress_report_author( |
| 54 | nickname=user_info.get("nickname"), |
| 55 | sec_uid=user_info.get("sec_uid") or sec_uid, |
| 56 | ) |
| 57 | |
| 58 | self._progress_update_step("下载模式", f"模式: {', '.join(modes)}") |
| 59 | |
| 60 | seen_aweme_ids: Set[str] = set() |
| 61 | for mode in modes: |
| 62 | strategy = self._get_mode_strategy(mode) |
| 63 | if strategy is None: |
| 64 | logger.warning("Unsupported user mode: %s", mode) |
| 65 | continue |
| 66 | |
| 67 | self._progress_update_step("下载模式", f"开始处理 {mode} 作品") |
| 68 | mode_result = await strategy.download_mode( |
| 69 | sec_uid, user_info, seen_aweme_ids=seen_aweme_ids |
| 70 | ) |
| 71 | result.total += mode_result.total |
| 72 | result.success += mode_result.success |
| 73 | result.failed += mode_result.failed |
| 74 | result.skipped += mode_result.skipped |
| 75 | |
| 76 | return result |
| 77 | |
| 78 | def _validate_mode_scope(self, sec_uid: str, modes: List[str]) -> bool: |
| 79 | normalized_modes = {str(mode or "").strip() for mode in modes} |
| 80 | has_collect_mode = bool(normalized_modes & self.SELF_COLLECT_MODES) |
| 81 | has_regular_mode = bool(normalized_modes - self.SELF_COLLECT_MODES) |
| 82 | |
| 83 | if has_collect_mode and sec_uid != "self": |
| 84 | # Desktop "我的内容 / 下载本收藏夹" sends the real self sec_uid |
| 85 | # together with a ``collects_id`` filter — by the time the |
| 86 | # request reaches here the sidecar has already verified via |
| 87 | # the cookie scope (``_resolve_viewer_sec_uid``) that the |
| 88 | # caller is the logged-in user, so a real sec_uid + collect |
| 89 | # mode + collects_id is the legit my-content path. Without |
| 90 | # this branch ``download()`` would short-circuit and produce |
| 91 | # an empty DownloadResult, which the JobManager renders as |
| 92 | # the silent "已完成 0 项" failure. |
| 93 | collects_id = (str(self.config.get("collects_id") or "")).strip() |
| 94 | if not collects_id: |
| 95 | logger.error( |
| 96 | "Modes collect/collectmix only support " |
| 97 | "/user/self?showTab=favorite_collection or " |
| 98 | "my-content 下载本收藏夹 (collects_id required)" |
| 99 | ) |
| 100 | return False |
| 101 | if has_collect_mode and has_regular_mode: |
| 102 | logger.error("Modes collect/collectmix cannot be combined with post/like/mix/music") |
| 103 | return False |
| 104 | return True |
| 105 | |
| 106 | def _filter_pinned_items(self, items: List[Dict[str, Any]]) -> List[Dict[str, Any]]: |
| 107 | if self._download_pinned_enabled(): |
| 108 | return items |
| 109 | return [item for item in items if not self._is_pinned_aweme(item)] |
| 110 | |
| 111 | def _download_pinned_enabled(self) -> bool: |
| 112 | return self._as_bool(self.config.get("download_pinned", False)) |
| 113 | |
| 114 | @staticmethod |
| 115 | def _is_pinned_aweme(item: Dict[str, Any]) -> bool: |
| 116 | value = item.get("is_top") |
| 117 | if isinstance(value, str): |
| 118 | return value.strip().lower() in {"1", "true", "yes", "on"} |
| 119 | return bool(value) |
| 120 | |
| 121 | @staticmethod |
| 122 | def _as_bool(value: Any) -> bool: |
| 123 | if isinstance(value, str): |
| 124 | return value.strip().lower() in {"1", "true", "yes", "on"} |
| 125 | return bool(value) |
| 126 | |
| 127 | async def _resolve_user_info(self, sec_uid: str, modes: List[str]) -> Optional[Dict[str, Any]]: |
| 128 | normalized_modes = {str(mode or "").strip() for mode in modes} |
| 129 | if sec_uid == "self" and normalized_modes.issubset(self.SELF_COLLECT_MODES): |
| 130 | self._progress_update_step("获取作者信息", "使用当前登录账号收藏夹上下文") |
| 131 | return { |
| 132 | "uid": "self", |
| 133 | "sec_uid": "self", |
| 134 | "nickname": "self", |
| 135 | } |
| 136 | |
| 137 | # Desktop my-content "下载本收藏夹" path: real sec_uid + collect |
| 138 | # mode + collects_id filter. The cookie scope upstream already |
| 139 | # guarantees this is the viewer themselves, so we can skip the |
| 140 | # network round-trip via ``api_client.get_user_info``. |
| 141 | if ( |
| 142 | normalized_modes.issubset(self.SELF_COLLECT_MODES) |
| 143 | and (str(self.config.get("collects_id") or "")).strip() |
| 144 | ): |
| 145 | self._progress_update_step("获取作者信息", "使用当前登录账号收藏夹上下文") |
| 146 | return { |
| 147 | "uid": sec_uid, |
| 148 | "sec_uid": sec_uid, |
| 149 | "nickname": "self", |
| 150 | } |
| 151 | |
| 152 | self._progress_update_step("获取作者信息", f"sec_uid={sec_uid}") |
| 153 | return await self.api_client.get_user_info(sec_uid) |
| 154 | |
| 155 | def _get_mode_strategy(self, mode: str): |
| 156 | normalized_mode = (mode or "").strip() |
| 157 | |
| 158 | # The "collect" strategy supports an optional ``collects_id`` filter |
| 159 | # that constrains paging to a single folder (desktop "我的收藏 / 下载 |
| 160 | # 本收藏夹"). When the filter is set we bypass the cache so the next |
| 161 | # call with a different (or absent) filter doesn't reuse a stale |
| 162 | # strategy bound to the previous folder. The no-filter path keeps |
| 163 | # caching to preserve the existing CLI behaviour. |
| 164 | if normalized_mode == "collect": |
| 165 | return self._make_collect_strategy() |
| 166 | |
| 167 | if normalized_mode in self._mode_strategy_cache: |
| 168 | return self._mode_strategy_cache[normalized_mode] |
| 169 | |
| 170 | strategy_cls = self.mode_registry.get(normalized_mode) |
| 171 | if strategy_cls is None: |
| 172 | return None |
| 173 | |
| 174 | strategy = strategy_cls(self) |
| 175 | self._mode_strategy_cache[normalized_mode] = strategy |
| 176 | return strategy |
| 177 | |
| 178 | def _make_collect_strategy(self): |
| 179 | """Construct the collect strategy, threading ``collects_id`` from |
| 180 | the per-job config when present. Caches only the no-filter path |
| 181 | (matching the historic CLI behaviour) so a subsequent call with a |
| 182 | different filter doesn't pick up a stale binding. |
| 183 | """ |
| 184 | strategy_cls = self.mode_registry.get("collect") |
| 185 | if strategy_cls is None: |
| 186 | return None |
| 187 | |
| 188 | raw_filter = self.config.get("collects_id") |
| 189 | collects_id = (str(raw_filter).strip() if raw_filter is not None else "") or None |
| 190 | |
| 191 | if collects_id is None: |
| 192 | cached = self._mode_strategy_cache.get("collect") |
| 193 | if cached is not None: |
| 194 | return cached |
| 195 | strategy = strategy_cls(self) |
| 196 | self._mode_strategy_cache["collect"] = strategy |
| 197 | return strategy |
| 198 | |
| 199 | # Filtered path is request-scoped — never cached. |
| 200 | return strategy_cls(self, collects_id=collects_id) |
| 201 | |
| 202 | async def _download_mode_items( |
| 203 | self, |
| 204 | mode: str, |
| 205 | items: List[Dict[str, Any]], |
| 206 | author_name: str, |
| 207 | seen_aweme_ids: Optional[Set[str]] = None, |
| 208 | ) -> DownloadResult: |
| 209 | if seen_aweme_ids is None: |
| 210 | seen_aweme_ids = set() |
| 211 | deduped_items: List[Dict[str, Any]] = [] |
| 212 | local_seen: Set[str] = set() |
| 213 | |
| 214 | for item in items: |
| 215 | aweme_id = str(item.get("aweme_id") or "").strip() |
| 216 | if not aweme_id: |
| 217 | continue |
| 218 | if aweme_id in seen_aweme_ids or aweme_id in local_seen: |
| 219 | continue |
| 220 | local_seen.add(aweme_id) |
| 221 | seen_aweme_ids.add(aweme_id) |
| 222 | deduped_items.append(item) |
| 223 | |
| 224 | result = DownloadResult() |
| 225 | result.total = len(deduped_items) |
| 226 | self._progress_set_item_total(result.total, "作品待下载") |
| 227 | self._progress_update_step("下载作品", f"待处理 {result.total} 条") |
| 228 | |
| 229 | # Accumulate per-aweme DB records and flush in a single transaction |
| 230 | # at the end — avoids one fsync per item across the whole batch. |
| 231 | db_batch: Optional[List[Dict[str, Any]]] = [] if self.database else None |
| 232 | |
| 233 | async def _process_aweme(item: Dict[str, Any]): |
| 234 | aweme_id = item.get("aweme_id") |
| 235 | if not await self._should_download(str(aweme_id or "")): |
| 236 | self._progress_advance_item("skipped", str(aweme_id or "unknown")) |
| 237 | return {"status": "skipped", "aweme_id": aweme_id} |
| 238 | |
| 239 | success = await self._download_aweme_assets( |
| 240 | item, author_name, mode=mode, db_batch=db_batch |
| 241 | ) |
| 242 | status = "success" if success else "failed" |
| 243 | self._progress_advance_item(status, str(aweme_id or "unknown")) |
| 244 | return { |
| 245 | "status": status, |
| 246 | "aweme_id": aweme_id, |
| 247 | } |
| 248 | |
| 249 | download_results = await self.queue_manager.download_batch(_process_aweme, deduped_items) |
| 250 | |
| 251 | if db_batch: |
| 252 | await self.database.add_aweme_batch(db_batch) |
| 253 | |
| 254 | for entry in download_results: |
| 255 | status = entry.get("status") if isinstance(entry, dict) else None |
| 256 | if status == "success": |
| 257 | result.success += 1 |
| 258 | elif status == "failed": |
| 259 | result.failed += 1 |
| 260 | elif status == "skipped": |
| 261 | result.skipped += 1 |
| 262 | else: |
| 263 | result.failed += 1 |
| 264 | self._progress_advance_item("failed", "unknown") |
| 265 | |
| 266 | return result |
| 267 | |
| 268 | # 向后兼容：旧测试仍直接调用 post 下载入口。 |
| 269 | async def _download_user_post(self, sec_uid: str, user_info: Dict[str, Any]) -> DownloadResult: |
| 270 | strategy = self._get_mode_strategy("post") |
| 271 | if strategy is None: |
| 272 | return DownloadResult() |
| 273 | return await strategy.download_mode(sec_uid, user_info, seen_aweme_ids=set()) |
| 274 | |
| 275 | async def _recover_user_post_with_browser( |
| 276 | self, |
| 277 | sec_uid: str, |
| 278 | user_info: Dict[str, Any], |
| 279 | aweme_list: List[Dict[str, Any]], |
| 280 | ) -> None: |
| 281 | browser_cfg = self.config.get("browser_fallback", {}) or {} |
| 282 | if not browser_cfg.get("enabled", True): |
| 283 | return |
| 284 | |
| 285 | number_limit = self.config.get("number", {}).get("post", 0) |
| 286 | # 在分页受限场景下，user_info.aweme_count 常常不可靠（经常只返回 20） |
| 287 | # 因此仅在用户显式设置 number_limit 时才限制浏览器采集目标数量。 |
| 288 | expected_count = int(number_limit or 0) |
| 289 | if expected_count and len(aweme_list) >= expected_count: |
| 290 | return |
| 291 | |
| 292 | try: |
| 293 | browser_aweme_ids = await self.api_client.collect_user_post_ids_via_browser( |
| 294 | sec_uid, |
| 295 | expected_count=expected_count, |
| 296 | headless=bool(browser_cfg.get("headless", False)), |
| 297 | max_scrolls=int(browser_cfg.get("max_scrolls", 240) or 240), |
| 298 | idle_rounds=int(browser_cfg.get("idle_rounds", 8) or 8), |
| 299 | wait_timeout_seconds=int(browser_cfg.get("wait_timeout_seconds", 600) or 600), |
| 300 | ) |
| 301 | except Exception as exc: |
| 302 | logger.error("Browser fallback failed: %s", exc) |
| 303 | return |
| 304 | |
| 305 | browser_aweme_items: Dict[str, Dict[str, Any]] = {} |
| 306 | browser_post_stats: Dict[str, int] = {} |
| 307 | if hasattr(self.api_client, "pop_browser_post_aweme_items"): |
| 308 | try: |
| 309 | browser_aweme_items = self.api_client.pop_browser_post_aweme_items() or {} |
| 310 | except Exception as exc: |
| 311 | logger.debug("Fetch browser post items skipped: %s", exc) |
| 312 | if hasattr(self.api_client, "pop_browser_post_stats"): |
| 313 | try: |
| 314 | browser_post_stats = self.api_client.pop_browser_post_stats() or {} |
| 315 | except Exception as exc: |
| 316 | logger.debug("Fetch browser post stats skipped: %s", exc) |
| 317 | |
| 318 | if not browser_aweme_ids: |
| 319 | logger.warning("Browser fallback returned no aweme_id") |
| 320 | return |
| 321 | |
| 322 | existing_ids = {str(item.get("aweme_id")) for item in aweme_list if item.get("aweme_id")} |
| 323 | missing_ids = [aweme_id for aweme_id in browser_aweme_ids if aweme_id not in existing_ids] |
| 324 | if not missing_ids: |
| 325 | return |
| 326 | |
| 327 | logger.warning( |
| 328 | "Recovering aweme details from browser list, missing count=%s", |
| 329 | len(missing_ids), |
| 330 | ) |
| 331 | detail_failed = 0 |
| 332 | detail_success = 0 |
| 333 | reused_from_browser_items = 0 |
| 334 | total_missing = len(missing_ids) |
| 335 | for index, aweme_id in enumerate(missing_ids, start=1): |
| 336 | if number_limit > 0 and len(aweme_list) >= number_limit: |
| 337 | break |
| 338 | |
| 339 | if index == 1 or index == total_missing or index % 5 == 0: |
| 340 | self._progress_update_step("浏览器回补", f"补全详情 {index}/{total_missing}") |
| 341 | |
| 342 | detail = browser_aweme_items.get(str(aweme_id)) |
| 343 | if not detail: |
| 344 | await self.rate_limiter.acquire() |
| 345 | detail = await self.api_client.get_video_detail(aweme_id, suppress_error=True) |
| 346 | if detail: |
| 347 | detail_success += 1 |
| 348 | else: |
| 349 | reused_from_browser_items += 1 |
| 350 | if not detail: |
| 351 | detail_failed += 1 |
| 352 | continue |
| 353 | author = detail.get("author", {}) if isinstance(detail, dict) else {} |
| 354 | detail_sec_uid = author.get("sec_uid") if isinstance(author, dict) else None |
| 355 | if detail_sec_uid and str(detail_sec_uid) != str(sec_uid): |
| 356 | logger.warning( |
| 357 | "Skip aweme_id=%s due to mismatched sec_uid (%s)", |
| 358 | aweme_id, |
| 359 | detail_sec_uid, |
| 360 | ) |
| 361 | continue |
| 362 | aweme_list.append(detail) |
| 363 | |
| 364 | self._progress_update_step( |
| 365 | "浏览器回补", |
| 366 | f"回补完成，复用 {reused_from_browser_items}，补拉成功 {detail_success}，失败 {detail_failed}", |
| 367 | ) |
| 368 | logger.warning( |
| 369 | "Browser fallback summary: merged_ids=%s selected_ids=%s post_items=%s post_pages=%s reused=%s detail_success=%s detail_failed=%s", |
| 370 | browser_post_stats.get("merged_ids", 0), |
| 371 | browser_post_stats.get("selected_ids", len(browser_aweme_ids)), |
| 372 | browser_post_stats.get("post_items", len(browser_aweme_items)), |
| 373 | browser_post_stats.get("post_pages", 0), |
| 374 | reused_from_browser_items, |
| 375 | detail_success, |
| 376 | detail_failed, |
| 377 | ) |
| 378 | |
| 379 | if detail_failed > 0: |
| 380 | logger.warning( |
| 381 | "Browser fallback detail fetch failed: %s/%s", |
| 382 | detail_failed, |
| 383 | total_missing, |
| 384 | ) |
| 385 |