| 1 | """抖音直播录制。 |
| 2 | |
| 3 | 技术路径： |
| 4 | - 通过 `/webcast/room/web/enter/` 获取 stream_url，常见字段： |
| 5 | * flv_pull_url: {SD, HD, FULL_HD, ORIGIN} |
| 6 | * hls_pull_url_map: {HD1, HD2, HD3} |
| 7 | - 选择最高清可用的流，优先 FLV（单文件落盘简单） |
| 8 | - 使用 aiohttp 分块写入到 `.flv` 临时文件，完成后原子重命名 |
| 9 | - 时长限制：read_timeout 自然结束或 max_duration_seconds 触发 |
| 10 | - 不依赖 ffmpeg；若用户需要转码可后处理 |
| 11 | |
| 12 | 限制： |
| 13 | - 不处理多人房间 / 连麦切换 |
| 14 | - 不采集弹幕（后续可扩展） |
| 15 | """ |
| 16 | |
| 17 | from __future__ import annotations |
| 18 | |
| 19 | import asyncio |
| 20 | import json |
| 21 | import os |
| 22 | import time |
| 23 | from datetime import datetime |
| 24 | from pathlib import Path |
| 25 | from typing import Any, Dict, Optional, Tuple |
| 26 | |
| 27 | import aiofiles |
| 28 | import aiohttp |
| 29 | |
| 30 | from core.downloader_base import BaseDownloader, DownloadResult |
| 31 | from utils.logger import setup_logger |
| 32 | from utils.naming import ( |
| 33 | DEFAULT_FILE_TEMPLATE, |
| 34 | DEFAULT_FOLDER_TEMPLATE, |
| 35 | build_live_context, |
| 36 | render_template, |
| 37 | ) |
| 38 | |
| 39 | logger = setup_logger("LiveDownloader") |
| 40 | |
| 41 | |
| 42 | # 质量优先级：数字越大越高清 |
| 43 | _FLV_QUALITY_ORDER = { |
| 44 | "ORIGIN": 100, |
| 45 | "FULL_HD1": 90, |
| 46 | "FULL_HD": 90, |
| 47 | "HD1": 70, |
| 48 | "HD": 70, |
| 49 | "SD1": 50, |
| 50 | "SD2": 50, |
| 51 | "SD": 50, |
| 52 | "LD": 30, |
| 53 | } |
| 54 | |
| 55 | |
| 56 | class LiveDownloader(BaseDownloader): |
| 57 | async def download(self, parsed_url: Dict[str, Any]) -> DownloadResult: |
| 58 | result = DownloadResult() |
| 59 | |
| 60 | room_id = parsed_url.get("room_id") |
| 61 | if not room_id: |
| 62 | logger.error("No room_id found in parsed URL") |
| 63 | return result |
| 64 | |
| 65 | result.total = 1 |
| 66 | self._progress_set_item_total(1, "直播录制") |
| 67 | self._progress_update_step("获取直播间信息", f"room_id={room_id}") |
| 68 | |
| 69 | info = await self.api_client.get_live_room_info(str(room_id)) |
| 70 | if not info: |
| 71 | logger.error("Live room not available or fetch failed: %s", room_id) |
| 72 | result.failed += 1 |
| 73 | self._progress_advance_item("failed", str(room_id)) |
| 74 | return result |
| 75 | |
| 76 | room = info.get("room") or {} |
| 77 | user = info.get("user") or {} |
| 78 | |
| 79 | status = room.get("status") |
| 80 | if status is not None and int(status or 0) != 2: |
| 81 | # 2 = 正在直播；其他状态不录 |
| 82 | logger.warning("Room %s not live (status=%s); skipping", room_id, status) |
| 83 | result.skipped += 1 |
| 84 | self._progress_advance_item("skipped", str(room_id)) |
| 85 | return result |
| 86 | |
| 87 | stream_url, quality = self._select_best_stream_url(room) |
| 88 | if not stream_url: |
| 89 | logger.error("No playable live stream URL for room %s", room_id) |
| 90 | result.failed += 1 |
| 91 | self._progress_advance_item("failed", str(room_id)) |
| 92 | return result |
| 93 | |
| 94 | author_name = (user.get("nickname") or "unknown").strip() or "unknown" |
| 95 | title = (room.get("title") or "直播").strip() or "直播" |
| 96 | save_dir, file_stem = self._plan_output_paths(author_name, title, str(room_id)) |
| 97 | |
| 98 | # 保存元数据 |
| 99 | meta_path = save_dir / f"{file_stem}_room.json" |
| 100 | try: |
| 101 | async with aiofiles.open(meta_path, "w", encoding="utf-8") as f: |
| 102 | await f.write(json.dumps(info, ensure_ascii=False, indent=2)) |
| 103 | except Exception as exc: |
| 104 | logger.debug("Save room meta failed: %s", exc) |
| 105 | |
| 106 | is_hls = ".m3u8" in stream_url.split("?")[0] |
| 107 | suffix = ".flv" if not is_hls else ".m3u8" |
| 108 | target_path = save_dir / f"{file_stem}{suffix}" |
| 109 | if is_hls: |
| 110 | # HLS 源只会下载 playlist（m3u8 文本），不是可直接播放的视频文件。 |
| 111 | # 告知用户正确的后处理方式。 |
| 112 | logger.warning( |
| 113 | "选中的直播源为 HLS（m3u8 playlist），保存的将是播放列表文本而非视频。" |
| 114 | "如需可播放文件，请用 ffmpeg 基于该 URL 抓流：ffmpeg -i '%s' -c copy out.ts", |
| 115 | stream_url, |
| 116 | ) |
| 117 | |
| 118 | live_cfg = self._live_config() |
| 119 | max_duration = float(live_cfg.get("max_duration_seconds") or 0) |
| 120 | chunk_size = int(live_cfg.get("chunk_size") or 65536) |
| 121 | idle_timeout = float(live_cfg.get("idle_timeout_seconds") or 30.0) |
| 122 | |
| 123 | self._progress_update_step( |
| 124 | "录制直播流", |
| 125 | f"quality={quality} | -> {target_path.name}", |
| 126 | ) |
| 127 | |
| 128 | ok = await self._record_stream( |
| 129 | stream_url, |
| 130 | target_path, |
| 131 | max_duration=max_duration, |
| 132 | chunk_size=chunk_size, |
| 133 | idle_timeout=idle_timeout, |
| 134 | ) |
| 135 | |
| 136 | if ok: |
| 137 | result.success += 1 |
| 138 | self._progress_advance_item("success", str(room_id)) |
| 139 | logger.info("Live recording finished: %s", target_path) |
| 140 | else: |
| 141 | result.failed += 1 |
| 142 | self._progress_advance_item("failed", str(room_id)) |
| 143 | |
| 144 | return result |
| 145 | |
| 146 | # --- helpers --- |
| 147 | |
| 148 | def _live_config(self) -> Dict[str, Any]: |
| 149 | cfg = self.config.get("live") or {} |
| 150 | return cfg if isinstance(cfg, dict) else {} |
| 151 | |
| 152 | def _plan_output_paths(self, author_name: str, title: str, room_id: str) -> Tuple[Path, str]: |
| 153 | started_at = datetime.now() |
| 154 | date = started_at.strftime("%Y-%m-%d_%H%M") |
| 155 | template_context = build_live_context( |
| 156 | room_id=str(room_id), |
| 157 | title=title, |
| 158 | author_name=author_name, |
| 159 | started_at=started_at, |
| 160 | ) |
| 161 | filename_template = self.config.get("filename_template") or DEFAULT_FILE_TEMPLATE |
| 162 | folder_template = self.config.get("folder_template") or DEFAULT_FOLDER_TEMPLATE |
| 163 | file_stem = render_template( |
| 164 | filename_template, |
| 165 | template_context, |
| 166 | fallback=f"{date}_{room_id}", |
| 167 | ) |
| 168 | folder_name = render_template( |
| 169 | folder_template, |
| 170 | template_context, |
| 171 | fallback=f"{date}_{room_id}", |
| 172 | ) |
| 173 | save_dir = self.file_manager.get_save_path( |
| 174 | author_name=author_name, |
| 175 | mode="live", |
| 176 | aweme_title=title, |
| 177 | aweme_id=room_id, |
| 178 | folderstyle=self.config.get("folderstyle", True), |
| 179 | download_date=date, |
| 180 | folder_name=folder_name, |
| 181 | author_sec_uid=None, |
| 182 | author_dir_style=self.config.get("author_dir") or "nickname", |
| 183 | ) |
| 184 | return save_dir, file_stem |
| 185 | |
| 186 | @staticmethod |
| 187 | def _select_best_stream_url(room: Dict[str, Any]) -> Tuple[Optional[str], str]: |
| 188 | """从 room.stream_url 中挑一条最佳地址。优先 FLV 高清。""" |
| 189 | stream = room.get("stream_url") if isinstance(room, dict) else None |
| 190 | if not isinstance(stream, dict): |
| 191 | return None, "" |
| 192 | |
| 193 | # FLV 优先 |
| 194 | flv_map = stream.get("flv_pull_url") |
| 195 | if isinstance(flv_map, dict) and flv_map: |
| 196 | best_key = max( |
| 197 | flv_map.keys(), |
| 198 | key=lambda k: _FLV_QUALITY_ORDER.get(k.upper(), 0), |
| 199 | ) |
| 200 | url = flv_map.get(best_key) |
| 201 | if isinstance(url, str) and url: |
| 202 | return url, best_key |
| 203 | |
| 204 | # 其次 HLS |
| 205 | hls_map = stream.get("hls_pull_url_map") |
| 206 | if isinstance(hls_map, dict) and hls_map: |
| 207 | best_key = max( |
| 208 | hls_map.keys(), |
| 209 | key=lambda k: _FLV_QUALITY_ORDER.get(k.upper(), 0), |
| 210 | ) |
| 211 | url = hls_map.get(best_key) |
| 212 | if isinstance(url, str) and url: |
| 213 | return url, best_key |
| 214 | |
| 215 | # 兜底：直接取根字段 |
| 216 | for key in ("flv_pull_url", "hls_pull_url", "rtmp_pull_url"): |
| 217 | url = stream.get(key) |
| 218 | if isinstance(url, str) and url: |
| 219 | return url, key |
| 220 | |
| 221 | return None, "" |
| 222 | |
| 223 | async def _record_stream( |
| 224 | self, |
| 225 | url: str, |
| 226 | target_path: Path, |
| 227 | *, |
| 228 | max_duration: float, |
| 229 | chunk_size: int, |
| 230 | idle_timeout: float, |
| 231 | ) -> bool: |
| 232 | """从 url 拉取字节流写入 target_path，直到流结束 / 超时 / 达到 max_duration。 |
| 233 | |
| 234 | **数据保留策略**：主播下播、网络空闲、payload 截断等场景下，只要已经写入 |
| 235 | > 0 字节，就把 .tmp 提升为正式文件（录到一半的直播也比零字节有用）。 |
| 236 | 仅 HTTP 4xx / 从未开始写入的情况下才会丢弃。 |
| 237 | """ |
| 238 | target_path.parent.mkdir(parents=True, exist_ok=True) |
| 239 | tmp_path = target_path.with_suffix(target_path.suffix + ".tmp") |
| 240 | start = time.monotonic() |
| 241 | bytes_written = 0 |
| 242 | last_chunk_ts = start |
| 243 | |
| 244 | # 直播 CDN 常同时校验 Referer 与 Origin 为 live.douyin.com（不是 www.douyin.com）。 |
| 245 | headers = self._download_headers() |
| 246 | headers["Referer"] = "https://live.douyin.com/" |
| 247 | headers["Origin"] = "https://live.douyin.com" |
| 248 | |
| 249 | def _promote_if_nonempty(reason: str) -> bool: |
| 250 | if bytes_written <= 0: |
| 251 | # 零字节也尝试清理 .tmp |
| 252 | try: |
| 253 | tmp_path.unlink(missing_ok=True) |
| 254 | except Exception: |
| 255 | pass |
| 256 | return False |
| 257 | try: |
| 258 | os.replace(str(tmp_path), str(target_path)) |
| 259 | except Exception as exc: |
| 260 | # 捕获所有异常：理论上只会是 OSError，但 rename 失败时宁可多兜底也别泄漏。 |
| 261 | logger.error("Live tmp → final rename failed: %s", exc) |
| 262 | return False |
| 263 | logger.info( |
| 264 | "Live stream recorded (%s): %s (%.1fs, %.1f MiB)", |
| 265 | reason, |
| 266 | target_path.name, |
| 267 | last_chunk_ts - start, |
| 268 | bytes_written / (1024 * 1024), |
| 269 | ) |
| 270 | return True |
| 271 | |
| 272 | session = await self.api_client.get_session() |
| 273 | try: |
| 274 | async with session.get( |
| 275 | url, |
| 276 | headers=headers, |
| 277 | timeout=aiohttp.ClientTimeout(total=None, sock_read=idle_timeout), |
| 278 | ) as resp: |
| 279 | if resp.status != 200: |
| 280 | logger.error("Live stream HTTP %s for %s", resp.status, target_path.name) |
| 281 | return False |
| 282 | async with aiofiles.open(tmp_path, "wb") as f: |
| 283 | async for chunk in resp.content.iter_chunked(chunk_size): |
| 284 | if not chunk: |
| 285 | continue |
| 286 | await f.write(chunk) |
| 287 | bytes_written += len(chunk) |
| 288 | now = time.monotonic() |
| 289 | last_chunk_ts = now |
| 290 | if max_duration and (now - start) >= max_duration: |
| 291 | logger.info( |
| 292 | "Live max_duration reached (%.1fs), stopping.", |
| 293 | max_duration, |
| 294 | ) |
| 295 | break |
| 296 | return _promote_if_nonempty("stream ended") |
| 297 | except asyncio.CancelledError: |
| 298 | # 外部取消（Ctrl+C 等）：保留已录制内容 |
| 299 | _promote_if_nonempty("cancelled") |
| 300 | raise |
| 301 | except aiohttp.ClientPayloadError as exc: |
| 302 | # 直播中断（主播下播）常见表现，视为正常结束 |
| 303 | logger.info("Live payload ended: %s", exc) |
| 304 | return _promote_if_nonempty("payload ended") |
| 305 | except (asyncio.TimeoutError, aiohttp.ServerTimeoutError) as exc: |
| 306 | # sock_read 空闲超时——多数情况是主播停止推流，保留已录数据 |
| 307 | logger.info("Live stream idle timeout after %ss: %s", idle_timeout, exc) |
| 308 | return _promote_if_nonempty("idle timeout") |
| 309 | except Exception as exc: |
| 310 | logger.error("Live stream recording failed: %s", exc) |
| 311 | # 其它未知异常也尽量保留已写入的数据 |
| 312 | return _promote_if_nonempty("unexpected error") |
| 313 |