| 1 | import argparse |
| 2 | import asyncio |
| 3 | import json |
| 4 | import re |
| 5 | import sys |
| 6 | from pathlib import Path |
| 7 | from typing import Any, Dict, List, Optional, Sequence |
| 8 | from urllib.parse import parse_qs, unquote, urlparse |
| 9 | |
| 10 | import yaml |
| 11 | |
| 12 | from utils.cookie_utils import parse_cookie_header, sanitize_cookies |
| 13 | |
| 14 | DEFAULT_URL = "https://www.douyin.com/" |
| 15 | DEFAULT_OUTPUT = Path("config/cookies.json") |
| 16 | REQUIRED_KEYS = {"msToken", "ttwid", "odin_tt", "passport_csrf_token"} |
| 17 | SUGGESTED_KEYS = REQUIRED_KEYS | {"sid_guard", "sessionid", "sid_tt"} |
| 18 | DEFAULT_AUXILIARY_KEYS = { |
| 19 | "_waftokenid", |
| 20 | "s_v_web_id", |
| 21 | "__ac_nonce", |
| 22 | "__ac_signature", |
| 23 | "UIFID", |
| 24 | "UIFID_TEMP", |
| 25 | "d_ticket", |
| 26 | "x-web-secsdk-uid", |
| 27 | "__security_server_data_status", |
| 28 | } |
| 29 | DEFAULT_AUXILIARY_PREFIXES = ( |
| 30 | "__security_mc_", |
| 31 | "bd_ticket_guard_", |
| 32 | "_bd_ticket_crypt_", |
| 33 | ) |
| 34 | PRIMARY_WAIT_UNTIL = "networkidle" |
| 35 | FALLBACK_WAIT_UNTIL = "domcontentloaded" |
| 36 | PRIMARY_TIMEOUT_MS = 300_000 |
| 37 | FALLBACK_TIMEOUT_MS = 300_000 |
| 38 | |
| 39 | |
| 40 | def parse_args(argv: Sequence[str]) -> argparse.Namespace: |
| 41 | parser = argparse.ArgumentParser( |
| 42 | description="Launch a browser, guide manual login, then dump Douyin cookies.", |
| 43 | ) |
| 44 | parser.add_argument( |
| 45 | "--url", |
| 46 | default=DEFAULT_URL, |
| 47 | help=f"Login page to open (default: {DEFAULT_URL})", |
| 48 | ) |
| 49 | parser.add_argument( |
| 50 | "--browser", |
| 51 | choices=["chromium", "firefox", "webkit"], |
| 52 | default="chromium", |
| 53 | help="Playwright browser engine (default: chromium)", |
| 54 | ) |
| 55 | parser.add_argument( |
| 56 | "--headless", |
| 57 | action="store_true", |
| 58 | help="Run browser headless (not recommended for manual login)", |
| 59 | ) |
| 60 | parser.add_argument( |
| 61 | "--output", |
| 62 | type=Path, |
| 63 | default=DEFAULT_OUTPUT, |
| 64 | help="JSON file to write collected cookies", |
| 65 | ) |
| 66 | parser.add_argument( |
| 67 | "--config", |
| 68 | type=Path, |
| 69 | help="Optional config.yml to update with captured cookies", |
| 70 | ) |
| 71 | parser.add_argument( |
| 72 | "--include-all", |
| 73 | action="store_true", |
| 74 | help="Store every cookie from douyin.com instead of the recommended subset", |
| 75 | ) |
| 76 | return parser.parse_args(argv) |
| 77 | |
| 78 | |
| 79 | async def capture_cookies(args: argparse.Namespace) -> int: |
| 80 | try: |
| 81 | from playwright.async_api import async_playwright # type: ignore |
| 82 | except ImportError: # pragma: no cover - defensive path |
| 83 | print( |
| 84 | "[ERROR] Playwright is not installed. Run `pip install playwright` first.", |
| 85 | file=sys.stderr, |
| 86 | ) |
| 87 | return 1 |
| 88 | |
| 89 | async with async_playwright() as p: |
| 90 | browser_factory = getattr(p, args.browser) |
| 91 | browser = await browser_factory.launch(headless=args.headless) |
| 92 | context = await browser.new_context() |
| 93 | page = await context.new_page() |
| 94 | observed_cookie_headers: List[str] = [] |
| 95 | observed_mstokens: List[str] = [] |
| 96 | |
| 97 | def _on_request(request: Any) -> None: |
| 98 | try: |
| 99 | headers = request.headers or {} |
| 100 | cookie_header = headers.get("cookie") |
| 101 | if cookie_header: |
| 102 | observed_cookie_headers.append(cookie_header) |
| 103 | url = request.url or "" |
| 104 | query = parse_qs(urlparse(url).query) |
| 105 | if "msToken" in query and query["msToken"]: |
| 106 | observed_mstokens.append((query["msToken"][0] or "").strip()) |
| 107 | token = extract_ms_token_from_text(url) |
| 108 | if token: |
| 109 | observed_mstokens.append(token) |
| 110 | except Exception: |
| 111 | # 观察请求失败不应影响主流程 |
| 112 | return |
| 113 | |
| 114 | page.on("request", _on_request) |
| 115 | |
| 116 | print("[INFO] Browser launched. Please complete Douyin login in the opened window.") |
| 117 | print("[INFO] Press Enter in this terminal once the homepage shows you are logged in.") |
| 118 | |
| 119 | await wait_for_login_confirmation(page, args.url) |
| 120 | |
| 121 | storage = await context.storage_state() |
| 122 | cookies = { |
| 123 | cookie["name"]: cookie["value"] |
| 124 | for cookie in storage["cookies"] |
| 125 | if cookie["domain"].endswith("douyin.com") |
| 126 | } |
| 127 | cookies = sanitize_cookies(cookies) |
| 128 | |
| 129 | ms_token = await try_extract_ms_token( |
| 130 | page, cookies, observed_cookie_headers, observed_mstokens |
| 131 | ) |
| 132 | if ms_token and not cookies.get("msToken"): |
| 133 | cookies["msToken"] = ms_token |
| 134 | print("[INFO] Extracted msToken from alternate sources.") |
| 135 | |
| 136 | await context.close() |
| 137 | await browser.close() |
| 138 | |
| 139 | picked = cookies if args.include_all else filter_cookies(cookies) |
| 140 | picked = sanitize_cookies(picked) |
| 141 | |
| 142 | args.output.parent.mkdir(parents=True, exist_ok=True) |
| 143 | args.output.write_text(json.dumps(picked, ensure_ascii=False, indent=2), encoding="utf-8") |
| 144 | print(f"[INFO] Saved {len(picked)} cookie(s) to {args.output.resolve()}") |
| 145 | |
| 146 | missing = REQUIRED_KEYS - picked.keys() |
| 147 | if missing: |
| 148 | print(f"[WARN] Missing required cookie keys: {', '.join(sorted(missing))}") |
| 149 | |
| 150 | if args.config: |
| 151 | update_config(args.config, picked) |
| 152 | |
| 153 | return 0 |
| 154 | |
| 155 | |
| 156 | def is_timeout_error(exc: Exception) -> bool: |
| 157 | return exc.__class__.__name__ == "TimeoutError" or "Timeout" in str(exc) |
| 158 | |
| 159 | |
| 160 | def is_target_closed_error(exc: Exception) -> bool: |
| 161 | return ( |
| 162 | exc.__class__.__name__ == "TargetClosedError" |
| 163 | or "Target page, context or browser has been closed" in str(exc) |
| 164 | ) |
| 165 | |
| 166 | |
| 167 | async def goto_with_fallback(page: Any, url: str) -> str: |
| 168 | # 部分站点会持续发请求，networkidle 可能一直达不到，超时后降级等待策略。 |
| 169 | try: |
| 170 | await page.goto(url, wait_until=PRIMARY_WAIT_UNTIL, timeout=PRIMARY_TIMEOUT_MS) |
| 171 | return PRIMARY_WAIT_UNTIL |
| 172 | except Exception as exc: |
| 173 | if is_target_closed_error(exc): |
| 174 | print( |
| 175 | "[WARN] Browser/page was closed during initial navigation, " |
| 176 | "continuing with current browser state." |
| 177 | ) |
| 178 | return "target_closed" |
| 179 | if not is_timeout_error(exc): |
| 180 | raise |
| 181 | print( |
| 182 | f"[WARN] goto(wait_until={PRIMARY_WAIT_UNTIL}) timed out after {PRIMARY_TIMEOUT_MS}ms, " |
| 183 | f"falling back to {FALLBACK_WAIT_UNTIL}." |
| 184 | ) |
| 185 | try: |
| 186 | await page.goto(url, wait_until=FALLBACK_WAIT_UNTIL, timeout=FALLBACK_TIMEOUT_MS) |
| 187 | return FALLBACK_WAIT_UNTIL |
| 188 | except Exception as exc: |
| 189 | if is_target_closed_error(exc): |
| 190 | print( |
| 191 | "[WARN] Browser/page was closed during fallback navigation, " |
| 192 | "continuing with current browser state." |
| 193 | ) |
| 194 | return "target_closed" |
| 195 | if is_timeout_error(exc): |
| 196 | print( |
| 197 | f"[WARN] goto(wait_until={FALLBACK_WAIT_UNTIL}) also timed out after {FALLBACK_TIMEOUT_MS}ms, " |
| 198 | "continuing anyway." |
| 199 | ) |
| 200 | return "timeout" |
| 201 | raise |
| 202 | |
| 203 | |
| 204 | async def wait_for_login_confirmation(page: Any, url: str, input_func: Any = input) -> None: |
| 205 | # 页面导航放到后台执行，避免在导航等待期间终端无法响应 Enter。 |
| 206 | nav_task = asyncio.create_task(goto_with_fallback(page, url)) |
| 207 | # 让 nav_task 至少进入第一个 await 点。否则在某些调度时序下， |
| 208 | # 若 input_func 立即返回（例如自动化测试或用户立刻按 Enter）， |
| 209 | # 可能导致 goto 尚未被调度便被 cancel，从而漏掉页面加载。 |
| 210 | await asyncio.sleep(0) |
| 211 | await asyncio.to_thread(input_func) |
| 212 | |
| 213 | if not nav_task.done(): |
| 214 | nav_task.cancel() |
| 215 | try: |
| 216 | await nav_task |
| 217 | except asyncio.CancelledError: |
| 218 | pass |
| 219 | except Exception as exc: |
| 220 | print(f"[WARN] Navigation task ended with error after cancel: {exc}") |
| 221 | return |
| 222 | |
| 223 | try: |
| 224 | await nav_task |
| 225 | except Exception as exc: |
| 226 | print(f"[WARN] Navigation task ended with error: {exc}") |
| 227 | |
| 228 | |
| 229 | async def try_extract_ms_token( |
| 230 | page: Any, |
| 231 | cookies: Dict[str, str], |
| 232 | observed_cookie_headers: List[str], |
| 233 | observed_mstokens: List[str], |
| 234 | ) -> Optional[str]: |
| 235 | existing = cookies.get("msToken") |
| 236 | if existing: |
| 237 | return existing |
| 238 | |
| 239 | for token in reversed(observed_mstokens): |
| 240 | token = (token or "").strip() |
| 241 | if token: |
| 242 | return token |
| 243 | |
| 244 | for header in reversed(observed_cookie_headers): |
| 245 | parsed = parse_cookie_header(header) |
| 246 | token = (parsed.get("msToken") or "").strip() |
| 247 | if token: |
| 248 | return token |
| 249 | extra = extract_ms_token_from_text(header) |
| 250 | if extra: |
| 251 | return extra |
| 252 | |
| 253 | try: |
| 254 | doc_cookie = await page.evaluate("() => document.cookie || ''") |
| 255 | parsed = parse_cookie_header(doc_cookie) |
| 256 | token = (parsed.get("msToken") or "").strip() |
| 257 | if token: |
| 258 | return token |
| 259 | extra = extract_ms_token_from_text(doc_cookie) |
| 260 | if extra: |
| 261 | return extra |
| 262 | except Exception: |
| 263 | pass |
| 264 | |
| 265 | js = """ |
| 266 | () => { |
| 267 | const values = []; |
| 268 | const pushIf = (v) => { |
| 269 | if (typeof v === 'string' && v.trim()) values.push(v.trim()); |
| 270 | }; |
| 271 | try { |
| 272 | for (const key of Object.keys(localStorage || {})) { |
| 273 | if (key.toLowerCase().includes('mstoken')) { |
| 274 | pushIf(localStorage.getItem(key)); |
| 275 | } |
| 276 | } |
| 277 | } catch (e) {} |
| 278 | try { |
| 279 | for (const key of Object.keys(sessionStorage || {})) { |
| 280 | if (key.toLowerCase().includes('mstoken')) { |
| 281 | pushIf(sessionStorage.getItem(key)); |
| 282 | } |
| 283 | } |
| 284 | } catch (e) {} |
| 285 | return values; |
| 286 | } |
| 287 | """ |
| 288 | try: |
| 289 | candidates = await page.evaluate(js) |
| 290 | for candidate in candidates or []: |
| 291 | if not isinstance(candidate, str): |
| 292 | continue |
| 293 | text = candidate.strip() |
| 294 | if not text: |
| 295 | continue |
| 296 | parsed = parse_cookie_header(text) |
| 297 | if parsed.get("msToken"): |
| 298 | return parsed["msToken"] |
| 299 | extra = extract_ms_token_from_text(text) |
| 300 | if extra: |
| 301 | return extra |
| 302 | if len(text) <= 2048 and all(ch not in text for ch in [";", " ", "\n", "\r", "\t"]): |
| 303 | return text |
| 304 | except Exception: |
| 305 | pass |
| 306 | |
| 307 | return None |
| 308 | |
| 309 | |
| 310 | def extract_ms_token_from_text(text: str) -> Optional[str]: |
| 311 | if not text: |
| 312 | return None |
| 313 | |
| 314 | patterns = [ |
| 315 | r"(?:^|[;,&\s\"'])msToken=([^;,&\s\"']+)", |
| 316 | r'"msToken"\s*:\s*"([^"]+)"', |
| 317 | r"'msToken'\s*:\s*'([^']+)'", |
| 318 | ] |
| 319 | for pattern in patterns: |
| 320 | match = re.search(pattern, text) |
| 321 | if not match: |
| 322 | continue |
| 323 | token = (match.group(1) or "").strip() |
| 324 | if token: |
| 325 | return unquote(token) |
| 326 | return None |
| 327 | |
| 328 | |
| 329 | def filter_cookies(cookies: Dict[str, str]) -> Dict[str, str]: |
| 330 | cookies = sanitize_cookies(cookies) |
| 331 | picked = {} |
| 332 | for key, value in cookies.items(): |
| 333 | if key in SUGGESTED_KEYS or key in DEFAULT_AUXILIARY_KEYS: |
| 334 | picked[key] = value |
| 335 | continue |
| 336 | if any(key.startswith(prefix) for prefix in DEFAULT_AUXILIARY_PREFIXES): |
| 337 | picked[key] = value |
| 338 | |
| 339 | if not picked: |
| 340 | return cookies |
| 341 | return picked |
| 342 | |
| 343 | |
| 344 | def update_config(config_path: Path, cookies: Dict[str, str]) -> None: |
| 345 | existing: Dict[str, object] = {} |
| 346 | if config_path.exists(): |
| 347 | existing = yaml.safe_load(config_path.read_text(encoding="utf-8")) or {} |
| 348 | |
| 349 | existing["cookies"] = cookies |
| 350 | |
| 351 | config_path.parent.mkdir(parents=True, exist_ok=True) |
| 352 | config_path.write_text( |
| 353 | yaml.safe_dump(existing, allow_unicode=True, sort_keys=False), |
| 354 | encoding="utf-8", |
| 355 | ) |
| 356 | print(f"[INFO] Updated config file: {config_path.resolve()}") |
| 357 | |
| 358 | |
| 359 | def main(argv: Optional[Sequence[str]] = None) -> int: |
| 360 | args = parse_args(argv or sys.argv[1:]) |
| 361 | return asyncio.run(capture_cookies(args)) |
| 362 | |
| 363 | |
| 364 | if __name__ == "__main__": |
| 365 | raise SystemExit(main()) |
| 366 |