442 lines
14 KiB
Python
442 lines
14 KiB
Python
import json
|
||
import os
|
||
import time
|
||
from datetime import datetime
|
||
from typing import Optional
|
||
|
||
SESSIONS_DIR = os.path.join(
|
||
os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
|
||
"sessions",
|
||
)
|
||
|
||
|
||
def ensure_sessions_dir():
|
||
os.makedirs(SESSIONS_DIR, exist_ok=True)
|
||
|
||
|
||
def get_cookie_path(account_id: int) -> str:
|
||
ensure_sessions_dir()
|
||
return os.path.join(SESSIONS_DIR, f"account_{account_id}.json")
|
||
|
||
|
||
def convert_dycred_to_storage_state(data: dict) -> dict:
|
||
import time
|
||
from urllib.parse import unquote
|
||
import re
|
||
|
||
# 1. 解析 Cookie
|
||
cookies = []
|
||
raw_cookie = data.get("cookie") or ""
|
||
now = int(time.time())
|
||
far_future = now + 60 * 60 * 24 * 180 # 180天过期
|
||
|
||
for pair in raw_cookie.split(";"):
|
||
pair = pair.strip()
|
||
if not pair or "=" not in pair:
|
||
continue
|
||
name, value = pair.split("=", 1)
|
||
name = name.strip()
|
||
value = value.strip()
|
||
if not name or any(c in name for c in "()[]{}'\"\n \t\\"):
|
||
continue
|
||
|
||
# 敏感且通常为 HttpOnly 的登录态保持 HttpOnly
|
||
is_httponly = name.lower() in ("sessionid", "sessionid_ss")
|
||
cookies.append({
|
||
"name": name,
|
||
"value": value,
|
||
"domain": ".douyin.com",
|
||
"path": "/",
|
||
"expires": far_future,
|
||
"httpOnly": is_httponly,
|
||
"secure": True,
|
||
"sameSite": "None"
|
||
})
|
||
|
||
# 2. 构建 localStorage 条目
|
||
local_storage = []
|
||
if data.get("keys"):
|
||
local_storage.append({"name": "security-sdk/s_sdk_crypt_sdk", "value": str(data["keys"])})
|
||
if data.get("web_protect"):
|
||
local_storage.append({"name": "security-sdk/s_sdk_sign_data_key/web_protect", "value": str(data["web_protect"])})
|
||
if data.get("sec_uid"):
|
||
local_storage.append({"name": "web_runtime_security_uid", "value": str(data["sec_uid"])})
|
||
|
||
# 补充 tea_cache_tokens
|
||
uid = str(data.get("unique_id") or data.get("my_uid") or "")
|
||
if uid:
|
||
local_storage.append({
|
||
"name": "tea_cache_tokens",
|
||
"value": json.dumps({"user_unique_id": uid, "web_id": uid}, ensure_ascii=False)
|
||
})
|
||
|
||
origins = []
|
||
if local_storage:
|
||
origins.append({
|
||
"origin": "https://www.douyin.com",
|
||
"localStorage": local_storage
|
||
})
|
||
|
||
# 3. 组装 storage_state
|
||
storage_state = {
|
||
"cookies": cookies,
|
||
"origins": origins,
|
||
"user_agent": data.get("ua") or "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"
|
||
}
|
||
|
||
# 4. 数字 UID 与 WebSocket 处理
|
||
if uid:
|
||
try:
|
||
storage_state["my_uid"] = int(uid)
|
||
except ValueError:
|
||
pass
|
||
|
||
ws_url = data.get("frontier_ws_url")
|
||
if not ws_url and data.get("ws_urls"):
|
||
for url in data["ws_urls"]:
|
||
if url:
|
||
ws_url = url
|
||
break
|
||
|
||
if ws_url:
|
||
storage_state["frontier_ws_url"] = ws_url
|
||
def _q(url_str: str, key: str) -> str:
|
||
m = re.search(rf"[?&]{re.escape(key)}=([^&\s]+)", url_str)
|
||
return unquote(m.group(1)) if m else ""
|
||
|
||
sdk_cert = _q(ws_url, "sdk_cert")
|
||
ts_sign = _q(ws_url, "ts_sign")
|
||
if sdk_cert:
|
||
storage_state["sdk_cert"] = sdk_cert
|
||
if ts_sign:
|
||
storage_state["ts_sign"] = ts_sign
|
||
|
||
return storage_state
|
||
|
||
|
||
def _parse_tea_from_ls(origins: list) -> tuple[str, str]:
|
||
"""从 origins.localStorage 提取 (my_uid, web_id)，www 站优先，跳过伪造 tea。"""
|
||
my_uid = ""
|
||
web_id = ""
|
||
ordered = sorted(origins or [], key=lambda o: 0 if "www.douyin.com" in (o.get("origin") or "") else 1)
|
||
for origin in ordered:
|
||
for entry in origin.get("localStorage") or []:
|
||
name = entry.get("name") or ""
|
||
if "tea_cache" not in name.lower():
|
||
continue
|
||
try:
|
||
parsed = json.loads(entry.get("value") or "{}")
|
||
except Exception:
|
||
continue
|
||
uid = str(parsed.get("user_unique_id") or "").strip()
|
||
wid = str(parsed.get("web_id") or "").strip()
|
||
if uid and wid and uid == wid and len(uid) > 12:
|
||
continue
|
||
if uid and uid.isdigit() and not my_uid:
|
||
my_uid = uid
|
||
if wid and wid.isdigit() and not web_id:
|
||
web_id = wid
|
||
if my_uid and web_id:
|
||
return my_uid, web_id or my_uid
|
||
return my_uid, web_id or my_uid
|
||
|
||
|
||
def _ws_device_id(url: str) -> str:
|
||
import re
|
||
from urllib.parse import unquote
|
||
m = re.search(r"[?&]device_id=([^&\s]+)", url or "")
|
||
return unquote(m.group(1)) if m else ""
|
||
|
||
|
||
def normalize_storage_state_for_im(data: dict) -> dict:
|
||
"""保存/导入前修正 UID、device_id；保留创作者私信 WS（aid=2906 + sdk_cert）。"""
|
||
if not isinstance(data, dict) or not isinstance(data.get("cookies"), list):
|
||
return data
|
||
data = dict(data)
|
||
origins = data.get("origins") or []
|
||
tea_uid, tea_web_id = _parse_tea_from_ls(origins)
|
||
|
||
ws_url = str(data.get("frontier_ws_url") or "")
|
||
is_creator_ws = "aid=2906" in ws_url and "sdk_cert=" in ws_url
|
||
|
||
if is_creator_ws:
|
||
ws_dev = _ws_device_id(ws_url)
|
||
if ws_dev and ws_dev.isdigit():
|
||
data["my_uid"] = int(ws_dev)
|
||
my_uid = ws_dev
|
||
else:
|
||
my_uid = str(data.get("my_uid") or tea_uid or "")
|
||
if my_uid.isdigit():
|
||
data["my_uid"] = int(my_uid)
|
||
else:
|
||
my_uid = ""
|
||
elif tea_uid:
|
||
data["my_uid"] = int(tea_uid)
|
||
my_uid = tea_uid
|
||
else:
|
||
my_uid = str(data.get("my_uid") or "")
|
||
if my_uid.isdigit():
|
||
data["my_uid"] = int(my_uid)
|
||
else:
|
||
my_uid = ""
|
||
|
||
if ws_url and my_uid and not is_creator_ws:
|
||
ws_dev = _ws_device_id(ws_url)
|
||
token = ""
|
||
try:
|
||
from urllib.parse import parse_qs, urlparse
|
||
token = parse_qs(urlparse(ws_url).query).get("token", [""])[0]
|
||
except Exception:
|
||
pass
|
||
looks_built = bool(token) and len(token) < 40
|
||
if ws_dev and ws_dev.isdigit() and ws_dev != my_uid and looks_built:
|
||
data.pop("frontier_ws_url", None)
|
||
data.pop("sdk_cert", None)
|
||
data.pop("ts_sign", None)
|
||
data.pop("frontier_ws_built", None)
|
||
|
||
for origin in origins:
|
||
ls = origin.get("localStorage") or []
|
||
fixed = False
|
||
for entry in ls:
|
||
if entry.get("name") == "web_runtime_security_uid":
|
||
fixed = True
|
||
val = str(entry.get("value") or "")
|
||
if not val.isdigit() and my_uid:
|
||
entry["value"] = my_uid
|
||
break
|
||
if not fixed and my_uid:
|
||
ls.append({"name": "web_runtime_security_uid", "value": my_uid})
|
||
origin["localStorage"] = ls
|
||
|
||
data["origins"] = origins
|
||
return data
|
||
|
||
|
||
def validate_cookie_json(cookie_data: str) -> dict:
|
||
cookie_data = cookie_data.strip()
|
||
if cookie_data.startswith("DYCRED1."):
|
||
import base64
|
||
try:
|
||
b64_part = cookie_data.split(".", 1)[1]
|
||
decoded = base64.b64decode(b64_part).decode("utf-8")
|
||
data = json.loads(decoded)
|
||
return normalize_storage_state_for_im(convert_dycred_to_storage_state(data))
|
||
except Exception as e:
|
||
raise ValueError(f"解析 DYCRED1 凭证密文失败: {e}")
|
||
|
||
data = json.loads(cookie_data)
|
||
if not isinstance(data, dict):
|
||
raise ValueError("Cookie 必须是 JSON 对象")
|
||
if "cookies" not in data or not isinstance(data.get("cookies"), list):
|
||
raise ValueError("Cookie 格式无效，需包含 cookies 字段（Playwright storage_state 格式）")
|
||
return normalize_storage_state_for_im(data)
|
||
|
||
|
||
def write_cookie_file(account_id: int, cookie_data: str) -> str:
|
||
data = validate_cookie_json(cookie_data)
|
||
path = get_cookie_path(account_id)
|
||
with open(path, "w", encoding="utf-8") as f:
|
||
json.dump(data, f, ensure_ascii=False, indent=2)
|
||
return path
|
||
|
||
|
||
def read_cookie_file(account_id: int) -> Optional[str]:
|
||
path = get_cookie_path(account_id)
|
||
if not os.path.exists(path) or os.path.getsize(path) == 0:
|
||
return None
|
||
with open(path, "r", encoding="utf-8") as f:
|
||
return f.read()
|
||
|
||
|
||
def clear_cookie_file(account_id: int):
|
||
path = get_cookie_path(account_id)
|
||
if os.path.exists(path):
|
||
os.remove(path)
|
||
|
||
|
||
DOUYIN_LOGIN_COOKIES = {
|
||
"sessionid",
|
||
"sessionid_ss",
|
||
"sid_tt",
|
||
"sid_tt_ss",
|
||
"uid_tt",
|
||
"uid_tt_ss",
|
||
"sid_guard",
|
||
"passport_auth_status",
|
||
"passport_auth_status_ss",
|
||
"login_status",
|
||
"odin_tt",
|
||
}
|
||
|
||
IM_TOKEN_COOKIES = ("sessionid", "sessionid_ss")
|
||
|
||
|
||
def _is_login_cookie(name: str) -> bool:
|
||
return name.lower() in DOUYIN_LOGIN_COOKIES
|
||
|
||
|
||
def _cookie_not_expired(cookie: dict, now: float) -> bool:
|
||
expires = cookie.get("expires")
|
||
if expires in (None, -1, 0):
|
||
return True
|
||
try:
|
||
return float(expires) > now
|
||
except (TypeError, ValueError):
|
||
return True
|
||
|
||
|
||
def analyze_cookie_data(data: Optional[dict]) -> dict:
|
||
"""静态分析 Cookie 是否具备登录凭证且未过期（不等同于服务端仍认可）"""
|
||
result = {
|
||
"has_cookie": False,
|
||
"cookie_valid": False,
|
||
"has_login_token": False,
|
||
"cookie_expired": False,
|
||
"cookie_count": 0,
|
||
"login_cookie_count": 0,
|
||
"expires_at": None,
|
||
"reason": "未保存 Cookie",
|
||
"has_sessionid": False,
|
||
}
|
||
if not data or not isinstance(data, dict):
|
||
return result
|
||
|
||
cookies = data.get("cookies", [])
|
||
if not isinstance(cookies, list) or not cookies:
|
||
result["reason"] = "Cookie 数据为空"
|
||
return result
|
||
|
||
result["has_cookie"] = True
|
||
result["cookie_count"] = len(cookies)
|
||
|
||
now = time.time()
|
||
login_cookies = [c for c in cookies if c.get("name") and _is_login_cookie(c["name"])]
|
||
result["login_cookie_count"] = len(login_cookies)
|
||
result["has_login_token"] = len(login_cookies) > 0
|
||
|
||
if not login_cookies:
|
||
result["reason"] = "缺少登录凭证（sessionid / sid_guard 等）"
|
||
return result
|
||
|
||
valid_login = [c for c in login_cookies if _cookie_not_expired(c, now)]
|
||
expired_login = [c for c in login_cookies if not _cookie_not_expired(c, now)]
|
||
|
||
if valid_login:
|
||
expires_values = [
|
||
float(c["expires"]) for c in valid_login
|
||
if c.get("expires") not in (None, -1, 0)
|
||
]
|
||
if expires_values:
|
||
earliest = min(expires_values)
|
||
result["expires_at"] = datetime.utcfromtimestamp(earliest).isoformat()
|
||
|
||
if not valid_login:
|
||
result["cookie_expired"] = True
|
||
result["reason"] = "登录 Cookie 已全部过期，需重新扫码"
|
||
return result
|
||
|
||
if expired_login:
|
||
result["reason"] = "部分登录 Cookie 已过期，启动后将探测浏览器是否已登录"
|
||
else:
|
||
has_sessionid = any(
|
||
c.get("name") in IM_TOKEN_COOKIES and c.get("value")
|
||
for c in cookies
|
||
)
|
||
if has_sessionid:
|
||
result["reason"] = "Cookie 有效，含 sessionid，可尝试 IM 直连"
|
||
else:
|
||
result["reason"] = "Cookie 有效，但缺少 sessionid，启动时将打开浏览器补全 IM 凭证"
|
||
|
||
result["cookie_valid"] = True
|
||
result["has_sessionid"] = any(
|
||
c.get("name") in IM_TOKEN_COOKIES and c.get("value") for c in cookies
|
||
)
|
||
return result
|
||
|
||
|
||
def analyze_cookie(cookie_data: Optional[str]) -> dict:
|
||
if not cookie_data:
|
||
return analyze_cookie_data(None)
|
||
try:
|
||
return analyze_cookie_data(json.loads(cookie_data))
|
||
except Exception:
|
||
return {
|
||
"has_cookie": False,
|
||
"cookie_valid": False,
|
||
"has_login_token": False,
|
||
"cookie_expired": False,
|
||
"cookie_count": 0,
|
||
"login_cookie_count": 0,
|
||
"expires_at": None,
|
||
"reason": "Cookie JSON 格式错误",
|
||
"has_sessionid": False,
|
||
}
|
||
|
||
|
||
def cookie_summary(cookie_data: Optional[str]) -> dict:
|
||
analysis = analyze_cookie(cookie_data)
|
||
key_names = []
|
||
if cookie_data:
|
||
try:
|
||
cookies = json.loads(cookie_data).get("cookies", [])
|
||
key_names = [c.get("name", "") for c in cookies if c.get("name")]
|
||
except Exception:
|
||
pass
|
||
return {
|
||
"cookie_count": analysis["cookie_count"],
|
||
"key_names": key_names[:20],
|
||
"cookie_valid": analysis["cookie_valid"],
|
||
"has_login_token": analysis["has_login_token"],
|
||
"cookie_expired": analysis["cookie_expired"],
|
||
"expires_at": analysis["expires_at"],
|
||
"reason": analysis["reason"],
|
||
"has_sessionid": analysis.get("has_sessionid", False),
|
||
}
|
||
|
||
|
||
def merge_playwright_cookies(storage: dict, live_cookies: list) -> dict:
|
||
"""将浏览器实时 Cookie 合并进 storage_state，确保 HttpOnly 的 sessionid 被保存"""
|
||
storage = dict(storage or {})
|
||
storage.setdefault("cookies", [])
|
||
storage.setdefault("origins", storage.get("origins") or [])
|
||
|
||
def cookie_key(c: dict) -> tuple:
|
||
return (c.get("name") or "", c.get("domain") or "", c.get("path") or "/")
|
||
|
||
merged: dict[tuple, dict] = {}
|
||
for item in storage.get("cookies") or []:
|
||
if item.get("name"):
|
||
merged[cookie_key(item)] = item
|
||
|
||
priority_names = set(IM_TOKEN_COOKIES) | DOUYIN_LOGIN_COOKIES
|
||
|
||
for item in live_cookies or []:
|
||
name = item.get("name") or ""
|
||
if not name:
|
||
continue
|
||
key = cookie_key(item)
|
||
entry = {
|
||
"name": name,
|
||
"value": item.get("value") or "",
|
||
"domain": item.get("domain") or "",
|
||
"path": item.get("path") or "/",
|
||
"expires": item.get("expires", -1),
|
||
"httpOnly": item.get("httpOnly", False),
|
||
"secure": item.get("secure", False),
|
||
"sameSite": item.get("sameSite", "Lax"),
|
||
}
|
||
prev = merged.get(key)
|
||
if not prev:
|
||
merged[key] = entry
|
||
continue
|
||
prev_val = prev.get("value") or ""
|
||
new_val = entry.get("value") or ""
|
||
if name in priority_names and len(new_val) > len(prev_val):
|
||
merged[key] = entry
|
||
elif not prev_val and new_val:
|
||
merged[key] = entry
|
||
|
||
storage["cookies"] = list(merged.values())
|
||
return storage
|