| 1 | from __future__ import annotations |
| 2 | |
| 3 | import json |
| 4 | import random |
| 5 | import string |
| 6 | import time |
| 7 | import urllib.request |
| 8 | from http.cookies import SimpleCookie |
| 9 | from threading import Lock |
| 10 | from typing import Any, Dict, Optional |
| 11 | |
| 12 | import yaml |
| 13 | |
| 14 | from utils.logger import setup_logger |
| 15 | |
| 16 | logger = setup_logger("MsTokenManager") |
| 17 | |
| 18 | |
| 19 | class MsTokenManager: |
| 20 | """ |
| 21 | 参考 F2 的 TokenManager 实现： |
| 22 | 1) 优先尝试从 mssdk 接口生成真实 msToken |
| 23 | 2) 失败时回退到随机 msToken，保证请求参数完整 |
| 24 | """ |
| 25 | |
| 26 | F2_CONF_URL = "https://raw.githubusercontent.com/Johnserf-Seed/f2/main/f2/conf/conf.yaml" |
| 27 | _cached_conf: Optional[Dict[str, Any]] = None |
| 28 | _cached_at: float = 0 |
| 29 | _cache_ttl_seconds: int = 3600 |
| 30 | _lock = Lock() |
| 31 | |
| 32 | def __init__( |
| 33 | self, |
| 34 | user_agent: str, |
| 35 | conf_url: Optional[str] = None, |
| 36 | timeout_seconds: int = 15, |
| 37 | ): |
| 38 | self.user_agent = user_agent |
| 39 | self.conf_url = conf_url or self.F2_CONF_URL |
| 40 | self.timeout_seconds = timeout_seconds |
| 41 | |
| 42 | @classmethod |
| 43 | def _is_valid_ms_token(cls, token: Optional[str]) -> bool: |
| 44 | if not token or not isinstance(token, str): |
| 45 | return False |
| 46 | # 与 F2 保持一致，长度通常为 164 或 184 |
| 47 | return len(token.strip()) in (164, 184) |
| 48 | |
| 49 | @classmethod |
| 50 | def gen_false_ms_token(cls) -> str: |
| 51 | token = ( |
| 52 | "".join(random.choice(string.ascii_letters + string.digits) for _ in range(182)) + "==" |
| 53 | ) |
| 54 | logger.debug("Generated fallback msToken") |
| 55 | return token |
| 56 | |
| 57 | def ensure_ms_token(self, cookies: Dict[str, str]) -> str: |
| 58 | current = (cookies or {}).get("msToken", "").strip() |
| 59 | if current: |
| 60 | return current |
| 61 | |
| 62 | real = self.gen_real_ms_token() |
| 63 | if real: |
| 64 | return real |
| 65 | |
| 66 | return self.gen_false_ms_token() |
| 67 | |
| 68 | def gen_real_ms_token(self) -> Optional[str]: |
| 69 | conf = self._load_f2_ms_token_conf() |
| 70 | if not conf: |
| 71 | return None |
| 72 | |
| 73 | payload = { |
| 74 | "magic": conf["magic"], |
| 75 | "version": conf["version"], |
| 76 | "dataType": conf["dataType"], |
| 77 | "strData": conf["strData"], |
| 78 | "ulr": conf["ulr"], |
| 79 | "tspFromClient": int(time.time() * 1000), |
| 80 | } |
| 81 | |
| 82 | request = urllib.request.Request( |
| 83 | conf["url"], |
| 84 | data=json.dumps(payload).encode("utf-8"), |
| 85 | headers={ |
| 86 | "Content-Type": "application/json; charset=utf-8", |
| 87 | "User-Agent": self.user_agent, |
| 88 | }, |
| 89 | method="POST", |
| 90 | ) |
| 91 | |
| 92 | try: |
| 93 | with urllib.request.urlopen(request, timeout=self.timeout_seconds) as resp: |
| 94 | token = self._extract_ms_token_from_headers(resp.headers) |
| 95 | if self._is_valid_ms_token(token): |
| 96 | logger.debug("Generated real msToken via mssdk endpoint") |
| 97 | return token |
| 98 | if token: |
| 99 | logger.warning("Generated msToken has unexpected length: %s", len(token.strip())) |
| 100 | return None |
| 101 | except Exception as exc: |
| 102 | logger.warning("Failed to generate real msToken: %s", exc) |
| 103 | return None |
| 104 | |
| 105 | def _load_f2_ms_token_conf(self) -> Optional[Dict[str, Any]]: |
| 106 | now = time.time() |
| 107 | with self._lock: |
| 108 | if self._cached_conf and (now - self._cached_at) < self._cache_ttl_seconds: |
| 109 | return self._cached_conf |
| 110 | |
| 111 | try: |
| 112 | with urllib.request.urlopen(self.conf_url, timeout=self.timeout_seconds) as resp: |
| 113 | raw = resp.read().decode("utf-8") |
| 114 | data = yaml.safe_load(raw) or {} |
| 115 | ms_conf = ( |
| 116 | data.get("f2", {}).get("douyin", {}).get("msToken", {}) # type: ignore[union-attr] |
| 117 | ) |
| 118 | |
| 119 | required = {"url", "magic", "version", "dataType", "ulr", "strData"} |
| 120 | if not required.issubset(ms_conf.keys()): |
| 121 | logger.warning( |
| 122 | "F2 msToken config incomplete, missing: %s", |
| 123 | sorted(required - set(ms_conf.keys())), |
| 124 | ) |
| 125 | return None |
| 126 | |
| 127 | with self._lock: |
| 128 | self._cached_conf = ms_conf |
| 129 | self._cached_at = now |
| 130 | return ms_conf |
| 131 | except Exception as exc: |
| 132 | logger.warning("Failed to load F2 msToken config: %s", exc) |
| 133 | return None |
| 134 | |
| 135 | @staticmethod |
| 136 | def _extract_ms_token_from_headers(headers: Any) -> Optional[str]: |
| 137 | set_cookies = headers.get_all("Set-Cookie") if hasattr(headers, "get_all") else [] |
| 138 | for header in set_cookies or []: |
| 139 | cookie = SimpleCookie() |
| 140 | cookie.load(header) |
| 141 | morsel = cookie.get("msToken") |
| 142 | if morsel and morsel.value: |
| 143 | return morsel.value.strip() |
| 144 | return None |
| 145 |