dy/backend/rpa_engine/douyin_im/frontier.py

"""Build frontier-im WebSocket URL (DouYin_Spider douyin_recv_msg logic)."""
import hashlib
import logging
import re
from typing import Optional
from urllib.parse import unquote

import requests

from .auth import DouyinAuth
from .dy_util import DEFAULT_USER_AGENT, generate_a_bogus, generate_msToken, generate_webid, splice_url
from .session import DouyinImSession, is_frontier_ws_url

logger = logging.getLogger("douyin_im.frontier")

APP_KEY = "e1bd35ec9db7b8d846de66ed140b1ad9"
FP_ID = "9"


def build_frontier_ws_url(session: DouyinImSession, device_id: str) -> Optional[str]:
    token = session.cookies.get("sessionid") or session.cookies.get("sessionid_ss") or ""
    if not token or not device_id:
        return None
    access_key_raw = f"{FP_ID}{APP_KEY}{device_id}f8a69f1719916z"
    access_key = hashlib.md5(access_key_raw.encode("utf-8")).hexdigest()
    params = {
        "aid": "6383",
        "device_platform": "douyin_pc",
        "fpid": FP_ID,
        "device_id": device_id,
        "token": token,
        "access_key": access_key,
    }
    query = "&".join(f"{k}={v}" for k, v in params.items())
    return f"wss://frontier-im.douyin.com/ws/v2?{query}"


def fetch_device_id(session: DouyinImSession) -> str:
    """Call Douyin user query API to obtain device/web id."""
    auth = DouyinAuth()
    auth.perepare_auth(
        session.cookie_header(),
        session.web_protect_str,
        session.keys_str,
        user_agent=session.user_agent or DEFAULT_USER_AGENT,
    )
    url = "https://www.douyin.com/aweme/v1/web/query/user"
    headers = {
        "User-Agent": session.user_agent,
        "Referer": "https://www.douyin.com/discover",
        "Accept": "application/json, text/plain, */*",
        "Cookie": session.cookie_header(),
    }
    params = {
        "device_platform": "webapp",
        "aid": "6383",
        "channel": "channel_pc_web",
        "publish_video_strategy_type": "2",
        "verifyFp": auth.cookie.get("s_v_web_id", "") if auth.cookie else "",
        "fp": auth.cookie.get("s_v_web_id", "") if auth.cookie else "",
        "webid": generate_webid(auth, "https://www.douyin.com/discover"),
        "msToken": generate_msToken(),
    }
    query = splice_url(params)
    params["a_bogus"] = generate_a_bogus(query, user_agent=session.user_agent)
    try:
        resp = requests.get(
            url,
            params=params,
            headers=headers,
            cookies=auth.cookie,
            verify=False,
            timeout=15,
        )
        data = resp.json()
        device_id = str(data.get("id") or data.get("device_id") or "")
        if device_id.isdigit():
            logger.info(f"Fetched device_id: {device_id[:20]}...")
            return device_id
        logger.warning(f"query/user returned non-numeric id: {device_id[:32]!r}")
    except Exception as e:
        logger.warning(f"fetch_device_id failed: {e}")
    return ""


def resolve_frontier_device_id(session: DouyinImSession) -> str:
    """Frontier WS requires numeric device_id from Douyin query/user API."""
    current = str(session.device_id or session.web_id or "")
    if current.isdigit():
        return current

    fetched = fetch_device_id(session)
    if fetched and str(fetched).isdigit():
        session.device_id = str(fetched)
        logger.info(f"Using numeric device_id for frontier WS: {fetched[:16]}...")
        return str(fetched)

    logger.warning(
        f"Invalid frontier device_id={current[:24]!r}; "
        "expected numeric id from query/user API"
    )
    return ""


def _ws_device_id(url: str) -> str:
    m = re.search(r"[?&]device_id=([^&\s]+)", url or "")
    return unquote(m.group(1)) if m else ""


def _ws_device_matches_session(session: DouyinImSession, url: str) -> bool:
    ws_dev = _ws_device_id(url)
    if not ws_dev or not ws_dev.isdigit():
        return True
    for candidate in (session.my_uid, session.web_id, session.device_id):
        if candidate and str(candidate) == ws_dev:
            return True
    return False


def _ws_token_looks_encoded(url: str) -> bool:
    m = re.search(r"[?&]token=([^&\s]+)", url or "")
    if not m:
        return False
    token = unquote(m.group(1))
    return len(token) >= 40 or not token.replace("_", "").replace("-", "").isalnum()


def ensure_frontier_ws(session: DouyinImSession) -> Optional[str]:
    """Ensure session has a usable frontier WebSocket URL."""
    session.sanitize_ws_urls()

    for url in session.ws_urls:
        if is_frontier_ws_url(url) and "sdk_cert=" in url and _ws_token_looks_encoded(url):
            session.ws_urls = [url]
            logger.info("Using captured real frontier WS URL (with sdk_cert)")
            return url

    # Only fpid=9 is Douyin private messaging.  Other Frontier products (for
    # example fpid=971 opened by the generic /chat page shell) also handshake
    # successfully but never carry this account's IM push stream.
    for url in session.ws_urls:
        if (
            is_frontier_ws_url(url)
            and "zijieapi.com" in url
            and "access_key=" in url
            and re.search(r"[?&]fpid=9(?:&|$)", url)
            and _ws_device_matches_session(session, url)
        ):
            session.ws_urls = [url]
            logger.info("Using captured browser frontier WebSocket URL")
            return url

    for url in session.ws_urls:
        if is_frontier_ws_url(url) and _ws_token_looks_encoded(url):
            session.ws_urls = [url]
            logger.info("Using captured frontier WS URL")
            return url

    # frontier 按 device_id 寻址推送，它不是账号 UID：抖音 query/user 返回的
    # id 才是本浏览器的设备注册号（my_uid 走 DouyinAuth，两者不能互换）。
    # 用 my_uid 拼出来的地址握手同样成功，但订阅的是另一个地址，
    # 于是长连接一直是「连上但收不到任何私信」。
    device_id = resolve_frontier_device_id(session)
    if not device_id:
        session.ws_urls = []
        return None

    built = build_frontier_ws_url(session, device_id)
    if built:
        session.ws_urls = [built]
        logger.info("Built frontier WS URL from cookie session")
        return built
    return None