227 lines
8.9 KiB
Python
227 lines
8.9 KiB
Python
"""账号凭证检测：静态 Cookie 分析 + IM 运行时校验"""
|
||
import asyncio
|
||
import json
|
||
import logging
|
||
from typing import Optional
|
||
|
||
from rpa_engine.douyin_im.auth import DouyinAuth
|
||
from rpa_engine.douyin_im.frontier import ensure_frontier_ws
|
||
from rpa_engine.douyin_im.http_client import DouyinImHttpClient
|
||
from rpa_engine.douyin_im.session import DouyinImSession
|
||
from utils.cookie_store import analyze_cookie
|
||
|
||
logger = logging.getLogger("credential")
|
||
|
||
|
||
def _should_reset_credentials(assessment: dict) -> bool:
|
||
"""凭证全面失效时需清空 Cookie/IM 数据并重新登录。"""
|
||
if not assessment.get("has_cookie"):
|
||
return False
|
||
if not assessment.get("cookie_valid"):
|
||
return True
|
||
message = assessment.get("message") or ""
|
||
# 仅缺 ticket/签名/浏览器采集 — 保留 Cookie，走浏览器补全即可
|
||
if any(
|
||
token in message
|
||
for token in ("ticket", "签名密钥", "浏览器模式", "web_protect")
|
||
):
|
||
return False
|
||
if (
|
||
not assessment.get("im_ready")
|
||
and not assessment.get("can_skip_browser")
|
||
and assessment.get("has_sessionid")
|
||
):
|
||
return True
|
||
return False
|
||
|
||
|
||
def build_im_session_from_storage(
|
||
storage: dict,
|
||
saved_im_data: Optional[str] = None,
|
||
) -> DouyinImSession:
|
||
session = DouyinImSession.from_storage_state(storage or {})
|
||
if saved_im_data:
|
||
try:
|
||
saved = DouyinImSession.from_dict(json.loads(saved_im_data))
|
||
# 新粘贴的 storage_state（含真实 frontier_ws_url）优先；只有它没带时才用缓存的。
|
||
if saved.ws_urls and not session.ws_urls:
|
||
session.ws_urls = saved.ws_urls
|
||
if saved.sdk_cert and not session.sdk_cert:
|
||
session.sdk_cert = saved.sdk_cert
|
||
if saved.frontier_ts_sign and not session.frontier_ts_sign:
|
||
session.frontier_ts_sign = saved.frontier_ts_sign
|
||
if saved.keys_str and not session.keys_str:
|
||
session.keys_str = saved.keys_str
|
||
if saved.web_protect_str and not session.web_protect_str:
|
||
session.web_protect_str = saved.web_protect_str
|
||
if saved.my_uid and not session.my_uid:
|
||
session.my_uid = saved.my_uid
|
||
if saved.device_id and not session.device_id:
|
||
session.device_id = saved.device_id
|
||
if saved.web_id and not session.web_id:
|
||
session.web_id = saved.web_id
|
||
if saved.conv_meta:
|
||
session.conv_meta = {**saved.conv_meta, **session.conv_meta}
|
||
except Exception:
|
||
pass
|
||
session.sanitize_ws_urls()
|
||
# Keep this builder pure and non-blocking. Frontier discovery can perform
|
||
# a synchronous network request (up to 15 seconds); callers that need it
|
||
# already do so from validate_im_session() through asyncio.to_thread() and
|
||
# the shared background-traffic limiter. Running it here made a large
|
||
# batch freeze the FastAPI event loop before any limiter was acquired.
|
||
return session
|
||
|
||
|
||
def has_im_session_token(session: DouyinImSession) -> bool:
|
||
return bool(session.cookies.get("sessionid") or session.cookies.get("sessionid_ss"))
|
||
|
||
|
||
def extract_sessionid_info(session: DouyinImSession) -> dict:
|
||
sessionid = session.cookies.get("sessionid") or ""
|
||
sessionid_ss = session.cookies.get("sessionid_ss") or ""
|
||
return {
|
||
"has_sessionid": bool(sessionid or sessionid_ss),
|
||
"sessionid": sessionid,
|
||
"sessionid_ss": sessionid_ss,
|
||
}
|
||
|
||
|
||
async def build_cookie_credential_detail(
|
||
cookie_data: Optional[str],
|
||
im_session_data: Optional[str] = None,
|
||
runtime_check: bool = True,
|
||
) -> dict:
|
||
"""供编辑账号页展示 IM 凭证与 sessionid 信息"""
|
||
sessionid_info = {
|
||
"has_sessionid": False,
|
||
"sessionid": "",
|
||
"sessionid_ss": "",
|
||
"im_ready": False,
|
||
"im_status": "未保存 Cookie",
|
||
"can_skip_browser": False,
|
||
"should_reset": False,
|
||
}
|
||
if not cookie_data:
|
||
return sessionid_info
|
||
|
||
try:
|
||
storage = json.loads(cookie_data)
|
||
session = build_im_session_from_storage(storage, im_session_data)
|
||
sessionid_info.update(extract_sessionid_info(session))
|
||
except Exception:
|
||
sessionid_info["im_status"] = "Cookie 格式错误"
|
||
return sessionid_info
|
||
|
||
if not runtime_check:
|
||
if sessionid_info["has_sessionid"]:
|
||
sessionid_info["im_status"] = "已检测到 sessionid（未做运行时验证）"
|
||
else:
|
||
sessionid_info["im_status"] = "缺少 sessionid，无法 IM 直连"
|
||
return sessionid_info
|
||
|
||
assessment = await assess_account_credential(cookie_data, im_session_data)
|
||
sessionid_info["im_ready"] = assessment["im_ready"]
|
||
sessionid_info["im_status"] = assessment["message"]
|
||
sessionid_info["can_skip_browser"] = assessment["can_skip_browser"]
|
||
sessionid_info["should_reset"] = assessment["should_reset"]
|
||
return sessionid_info
|
||
|
||
|
||
async def validate_im_session(
|
||
session: DouyinImSession,
|
||
_bypass_global_limit: bool = False,
|
||
) -> tuple[bool, str]:
|
||
if not _bypass_global_limit:
|
||
from rpa_engine.douyin_im.traffic_control import get_traffic_controller
|
||
|
||
controller = get_traffic_controller()
|
||
async with controller.background_slot(0, "credential validation"):
|
||
return await validate_im_session(session, _bypass_global_limit=True)
|
||
|
||
if not session.can_direct_im():
|
||
if not has_im_session_token(session):
|
||
return False, "缺少 sessionid，无法直连 IM"
|
||
return False, "Cookie 不满足 IM 直连条件"
|
||
|
||
await asyncio.to_thread(ensure_frontier_ws, session)
|
||
try:
|
||
auth = DouyinAuth.from_im_session(session)
|
||
# 优先用已持久化的 my_uid，避免每次都发起网络 query_my_uid（uid_tt 是加密串，
|
||
# int() 解析必然失败而回退到网络请求；该请求偶发失败会误判为“未就绪”）。
|
||
uid = session.my_uid or auth.get_uid()
|
||
if not uid:
|
||
return False, "服务端未认可当前 Cookie（无法获取用户 UID）"
|
||
if not auth.is_sign_ready():
|
||
return False, "缺少 IM 签名密钥（web_protect/keys），请用浏览器登录补全"
|
||
session.my_uid = int(uid)
|
||
async with DouyinImHttpClient(session) as http:
|
||
await http.get_unread_count()
|
||
# 若已缓存到会话票据，优先校验其是否仍新鲜（最理想）。
|
||
if session.conv_meta:
|
||
ok, reason = await http.verify_messaging_capability(auth, session.my_uid)
|
||
if ok:
|
||
return True, reason
|
||
# 没有缓存会话票据是首次登录的正常情况：会话 ticket 会在发送时即时
|
||
# 创建/获取（resolve_conversation_meta），因此只要 Cookie + sessionid +
|
||
# 签名密钥(web_protect/keys) + UID 齐全，就视为可 IM 直连托管，不必再开浏览器。
|
||
return True, "IM 凭证就绪（Cookie 与签名密钥齐全，可直连托管）"
|
||
except Exception as e:
|
||
logger.warning(f"IM session validation failed: {e}")
|
||
return False, f"IM 运行时验证失败: {e}"
|
||
|
||
|
||
async def assess_account_credential(
|
||
cookie_data: Optional[str],
|
||
im_session_data: Optional[str] = None,
|
||
) -> dict:
|
||
cookie_info = analyze_cookie(cookie_data)
|
||
result = {
|
||
"has_cookie": cookie_info.get("has_cookie", False),
|
||
"cookie_valid": cookie_info.get("cookie_valid", False),
|
||
"cookie_status": cookie_info.get("reason", ""),
|
||
"has_sessionid": False,
|
||
"im_ready": False,
|
||
"can_skip_browser": False,
|
||
"login_mode": "browser",
|
||
"message": "未保存 Cookie，需浏览器扫码登录",
|
||
"should_reset": False,
|
||
}
|
||
|
||
if not cookie_data:
|
||
return result
|
||
|
||
try:
|
||
storage = json.loads(cookie_data)
|
||
except Exception:
|
||
result["message"] = "Cookie 格式错误"
|
||
result["should_reset"] = _should_reset_credentials(result)
|
||
return result
|
||
|
||
session = build_im_session_from_storage(storage, im_session_data)
|
||
result["has_sessionid"] = has_im_session_token(session)
|
||
|
||
if not cookie_info.get("cookie_valid"):
|
||
result["message"] = cookie_info.get("reason") or "Cookie 无效，需重新登录"
|
||
result["should_reset"] = _should_reset_credentials(result)
|
||
return result
|
||
|
||
if not result["has_sessionid"]:
|
||
result["login_mode"] = "browser"
|
||
result["message"] = "Cookie 已保存但缺少 sessionid，需浏览器刷新登录态"
|
||
result["should_reset"] = _should_reset_credentials(result)
|
||
return result
|
||
|
||
im_ok, im_reason = await validate_im_session(session)
|
||
result["im_ready"] = im_ok
|
||
if im_ok:
|
||
result["can_skip_browser"] = True
|
||
result["login_mode"] = "im_direct"
|
||
result["message"] = im_reason
|
||
else:
|
||
result["login_mode"] = "browser"
|
||
result["message"] = im_reason
|
||
|
||
result["should_reset"] = _should_reset_credentials(result)
|
||
return result
|