spider_work/jhs_rpc_spider/jhs_raw_codec_client.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Parameter-based client for jhs_raw_codec_rpc.js

Usage (import-based):
    from raw_codec_rpc.jhs_raw_codec_client import call_codec

    enc = call_codec({
        "op": "enc",
        "url": "https://api.jihuanshe.com/api/market/auction-products?sorting=completed&page=2&token=..."
    })

    dec = call_codec({
        "op": "dec",
        "request_url": "https://api.jihuanshe.com/api/market/banners?raw_data=...&token=...",
        "response_raw_data": "BASE64_CIPHER"
    })
"""

from pathlib import Path
from typing import Any, Dict, Optional
import time
import json
import os
import queue
import subprocess
import tempfile
import threading
import concurrent.futures
from urllib.parse import unquote

import frida


PKG = "com.jihuanshe"
SCRIPT_PATH = Path(__file__).with_name("jhs_raw_codec_rpc.js")
ENV_DEVICE_ID = "FRIDA_DEVICE_ID"
ENV_DEBUG = "JHS_CODEC_DEBUG"
ENV_CLI_TARGET_SEC = "FRIDA_CLI_TARGET_SEC"
ENV_RPC_TIMEOUT_SEC = "FRIDA_RPC_TIMEOUT_SEC"


class JhsRawCodecClient:
    RETRYABLE_FRIDA_ERROR_MARKERS = (
        "unable to connect to remote frida-server: closed",
        "server is not running",
        "the connection is closed",
        "connection is closed",
        "connection terminated",
        "session is detached",
        "script is destroyed",
        "transport error",
        "device lost",
    )

    def _on_script_message(self, message, data) -> None:
        msg_type = message.get("type")
        if msg_type == "log":
            level = message.get("level", "log")
            payload = message.get("payload", "")
            print(f"[frida:{level}] {payload}")
            return
        if msg_type == "error":
            desc = message.get("description", "script error")
            stack = message.get("stack")
            print(f"[frida:error] {desc}")
            if stack:
                print(stack)
            return
        print(f"[frida:{msg_type}] {message}")

    def _resolve_device(self, device_id: Optional[str]):
        """
        解析并返回 Frida 设备对象。

        优先使用显式传入的 device_id；若为空则回退到 USB 设备。

        Args:
            device_id: 目标设备 ID，例如 emulator-5554。

        Returns:
            frida.core.Device: 已连接设备对象。
        """
        if device_id:
            return frida.get_device_manager().get_device(device_id, timeout=5)
        return frida.get_usb_device(timeout=5)

    def _log(self, msg: str) -> None:
        if self.debug:
            print(f"[JhsRawCodecClient] {msg}")

    def _is_retryable_frida_error(self, exc: Exception) -> bool:
        msg = str(exc).strip().lower()
        return any(marker in msg for marker in self.RETRYABLE_FRIDA_ERROR_MARKERS)

    def _find_pid_by_identifier(self) -> int:
        """
        通过应用标识符(package name)查找目标应用的进程 PID。

        Returns:
            int: 找到则返回 PID，未找到返回 0。
        """
        # Prefer app identifier lookup; attach("name") in frida-python matches process name.
        try:
            for app in self.device.enumerate_applications():
                if app.identifier == self.package and app.pid:
                    return int(app.pid)
        except Exception:
            pass
        return 0

    def _find_pid_by_process(self) -> int:
        """
        通过进程名查找目标进程 PID（作为 identifier 查找的兜底方案）。

        Returns:
            int: 找到则返回 PID，未找到返回 0。
        """
        try:
            for p in self.device.enumerate_processes():
                if p.name == self.package:
                    return int(p.pid)
        except Exception:
            pass
        return 0

    def _attach_session(self) -> None:
        last_err = None
        self.device = self._resolve_device(self.device_id)
        for _ in range(6):
            try:
                pid = self._find_pid_by_identifier() or self._find_pid_by_process()
                if pid:
                    self.session = self.device.attach(pid)
                    return
                last_err = frida.ProcessNotFoundError(
                    f"unable to find running app/process for '{self.package}'"
                )
                time.sleep(1.0)
            except frida.ProcessNotFoundError as e:
                last_err = e
                time.sleep(1.0)
        raise RuntimeError(
            f"unable to attach '{self.package}', please open app and keep it running"
        ) from last_err

    def _load_script(self) -> None:
        code = SCRIPT_PATH.read_text(encoding="utf-8")
        self.script = self.session.create_script(code)
        self.script.on("message", self._on_script_message)
        self.script.load()

    def _connect(self) -> None:
        self.session = None
        self.script = None
        self._attach_session()
        self._load_script()

    def reconnect(self) -> None:
        self._log("reconnecting frida session")
        self.close()
        self._connect()

    def __init__(
        self,
        package: str = PKG,
        device_id: Optional[str] = None,
        cli_target_sec: Optional[int] = None,
        rpc_timeout_sec: Optional[int] = None,
    ):
        """
        初始化客户端并附加到目标 App 进程，随后加载 RPC 脚本。

        Args:
            package: 目标应用包名，默认使用常量 PKG。
            device_id: Frida 设备 ID。未传时读取环境变量 FRIDA_DEVICE_ID，
                若仍为空则使用默认 USB 设备选择逻辑。
            cli_target_sec: CLI 兜底模式的 frida `-t` 秒数。未传时读取
                环境变量 FRIDA_CLI_TARGET_SEC，默认 3 秒。
            rpc_timeout_sec: Python RPC 同步调用的超时秒数（frida-python 的
                `exports_sync` 本身无超时，靠外层 future 兜底）。未传时读取
                环境变量 FRIDA_RPC_TIMEOUT_SEC，默认 20 秒；最小 3 秒。

        Raises:
            RuntimeError: 多次重试后仍无法附加到目标进程。
        """
        self.package = package
        self.device_id = device_id or os.getenv(ENV_DEVICE_ID)
        self.debug = os.getenv(ENV_DEBUG, "").strip().lower() in {"1", "true", "yes", "on"}
        target_sec = cli_target_sec
        if target_sec is None:
            target_sec = int(os.getenv(ENV_CLI_TARGET_SEC, "3"))
        self.cli_target_sec = max(1, int(target_sec))
        rpc_ts = rpc_timeout_sec
        if rpc_ts is None:
            rpc_ts = int(os.getenv(ENV_RPC_TIMEOUT_SEC, "20"))
        self.rpc_timeout_sec = max(3, int(rpc_ts))
        self.device = None
        self.session = None
        self.script = None
        self._prefer_cli = False
        # 单线程 executor 专门用来给同步 RPC 加超时；到点若卡死则丢弃换新的
        self._executor = concurrent.futures.ThreadPoolExecutor(
            max_workers=1, thread_name_prefix="jhs-rpc"
        )
        self._connect()

    def __enter__(self):
        """上下文管理器入口，返回当前客户端实例。"""
        return self

    def __exit__(self, exc_type, exc, tb):
        """上下文管理器退出时释放 Frida 资源。"""
        self.close()
        return False

    def close(self) -> None:
        """关闭并清理资源：卸载脚本、断开会话连接、关闭 executor。"""
        try:
            if self.script is not None:
                self.script.unload()
        except Exception:
            pass
        try:
            if self.session is not None:
                self.session.detach()
        except Exception:
            pass
        self.script = None
        self.session = None
        # executor 里可能残留卡死的调用线程，不 join，让守护线程随进程结束
        try:
            self._executor.shutdown(wait=False, cancel_futures=True)
        except Exception:
            pass

    def _sync_rpc_call(self, params: Dict[str, Any]) -> Dict[str, Any]:
        """
        为 `script.exports_sync.call` 增加超时保护的包装。

        frida-python 的 `exports_sync` 是纯阻塞调用、没有 timeout 参数；
        目标 app 冷启动时相关 Java 类可能还未加载，脚本 hook 排队后不返回，
        Python 端会永远阻塞。这里通过后台 future + `result(timeout=...)` 兜底。

        Args:
            params (Dict[str, Any]): RPC 参数字典。

        Returns:
            Dict[str, Any]: JS 侧返回的结果。

        Raises:
            TimeoutError: 超过 `self.rpc_timeout_sec` 仍未返回；抛出前会尝试
                detach session 让底层阻塞线程报错退出，并重建 executor。
        """
        fut = self._executor.submit(self.script.exports_sync.call, params)
        try:
            return fut.result(timeout=self.rpc_timeout_sec)
        except concurrent.futures.TimeoutError:
            self._log(f"python rpc timeout after {self.rpc_timeout_sec}s, detaching session")
            # 强断会话，通常会让阻塞的 frida 线程抛异常退出（否则它就常驻卡着）
            try:
                if self.session is not None:
                    self.session.detach()
            except Exception:
                pass
            # 旧 executor 内的线程可能还在阻塞，直接丢弃、开新的
            try:
                self._executor.shutdown(wait=False, cancel_futures=True)
            except Exception:
                pass
            self._executor = concurrent.futures.ThreadPoolExecutor(
                max_workers=1, thread_name_prefix="jhs-rpc"
            )
            raise TimeoutError(
                f"frida python rpc call timeout after {self.rpc_timeout_sec}s"
            )

    def encrypt(self, url: str) -> Dict[str, Any]:
        """
        调用 JS RPC 的 encrypt 方法，对请求 URL 进行加密处理。

        Args:
            url: 原始请求 URL。

        Returns:
            Dict[str, Any]: JS 侧返回的加密结果字典。
        """
        return self.script.exports_sync.encrypt(url)

    def decrypt(self, request_url_with_raw_data: str, response_raw_data: str) -> Dict[str, Any]:
        """
        调用 JS RPC 的 decrypt 方法，对响应中的 raw_data 进行解密。

        Args:
            request_url_with_raw_data: 包含 raw_data 参数的请求 URL。
            response_raw_data: 响应中的加密 raw_data 字符串。

        Returns:
            Dict[str, Any]: JS 侧返回的解密结果字典。
        """
        return self.script.exports_sync.decrypt(request_url_with_raw_data, response_raw_data)

    def call(self, params: Dict[str, Any]) -> Dict[str, Any]:
        """
        统一调用入口，转发到 JS RPC 的 call 方法。

        当 Python 会话环境缺失 Java bridge（如部分 Gadget 场景）时，
        自动降级为 CLI 注入方式调用。

        Args:
            params: RPC 调用参数字典。

        Returns:
            Dict[str, Any]: RPC 返回结果。
        """
        if self._prefer_cli:
            self._log("call path: cli(preferred)")
            return self._call_via_cli(params)
        try:
            self._log("call path: python rpc")
            return self._sync_rpc_call(params)
        except TimeoutError as e:
            # 已经在 _sync_rpc_call 内 detach 过；这里直接重连，让上层重试
            self._log(f"python rpc timeout, reconnecting session: {e}")
            try:
                self.reconnect()
            except Exception as re:
                self._log(f"reconnect after timeout failed: {re}")
            raise
        except Exception as e:
            msg = str(e)
            # In some embedded Gadget setups, Python session scripts miss Java bridge.
            if "Java is not defined" in msg or "ReferenceError: 'Java' is not defined" in msg:
                # Once detected, skip the slow exception path on later calls.
                self._prefer_cli = True
                self._log("python rpc missing Java bridge, switch to cli fallback")
                return self._call_via_cli(params)
            if self._is_retryable_frida_error(e):
                self._log(f"python rpc failed with retryable frida error: {e}")
                self.reconnect()
                return self._sync_rpc_call(params)
            raise

    def _call_via_cli(self, params: Dict[str, Any]) -> Dict[str, Any]:
        """
        使用 frida CLI 作为兜底方案执行一次 RPC 调用。

        实现流程：
        1) 拼接临时 JS（注入参数并调用 rpc.exports.call）
        2) 通过 frida 命令行注入目标进程执行
        3) 从标准输出解析约定的结果前缀

        Args:
            params: RPC 调用参数字典。

        Returns:
            Dict[str, Any]: RPC 返回结果。

        Raises:
            RuntimeError: CLI 调用失败或未解析到结果行。
        """
        js_src = SCRIPT_PATH.read_text(encoding="utf-8")
        params_json = json.dumps(params, ensure_ascii=False)
        wrapper = (
            "const __PARAMS = " + params_json + ";\n"
            + js_src
            + "\nsetImmediate(function(){\n"
              "  rpc.exports.call(__PARAMS)\n"
              "    .then(function(r){ console.log('[CODEC-RESULT-URI]' + encodeURIComponent(JSON.stringify(r))); })\n"
              "    .catch(function(e){ console.log('[CODEC-ERROR]' + e); });\n"
              "});\n"
        )
        fd, tmp_path = tempfile.mkstemp(prefix="jhs_codec_", suffix=".js")
        os.close(fd)
        proc = None
        try:
            Path(tmp_path).write_text(wrapper, encoding="utf-8")
            cli_cmd = ["frida"]
            if self.device_id:
                cli_cmd.extend(["-D", self.device_id])
            else:
                cli_cmd.append("-U")
            cli_cmd.extend(["-N", self.package, "-l", tmp_path, "-q", "-t", str(self.cli_target_sec)])
            self._log("spawn cli: " + " ".join(cli_cmd))
            proc = subprocess.Popen(
                cli_cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                encoding="utf-8",
                errors="replace",
            )
            deadline_total = max(10, self.cli_target_sec + 8)
            deadline = time.time() + deadline_total
            out_lines = []
            err_lines = []

            # 后台线程把 stdout / stderr 读到 queue 里，主循环用 queue.get(timeout=...)
            # 拿数据，从根本上避免 readline() 阻塞导致 deadline 失效
            io_q: "queue.Queue" = queue.Queue()

            def _pump(stream, tag):
                try:
                    for line in iter(stream.readline, ""):
                        io_q.put((tag, line))
                finally:
                    try:
                        stream.close()
                    except Exception:
                        pass
                    io_q.put((tag, None))  # EOF 标记

            t_out = threading.Thread(target=_pump, args=(proc.stdout, "out"), daemon=True)
            t_err = threading.Thread(target=_pump, args=(proc.stderr, "err"), daemon=True)
            t_out.start()
            t_err.start()

            out_eof = False
            err_eof = False

            while True:
                remaining = deadline - time.time()
                if remaining <= 0:
                    break
                try:
                    tag, line = io_q.get(timeout=min(0.5, remaining))
                except queue.Empty:
                    # 进程已结束且两路输出都到 EOF 就退出，否则继续等
                    if proc.poll() is not None and out_eof and err_eof:
                        break
                    continue

                if line is None:
                    if tag == "out":
                        out_eof = True
                    else:
                        err_eof = True
                    if out_eof and err_eof and proc.poll() is not None:
                        break
                    continue

                line = line.rstrip("\r\n")
                if tag == "err":
                    err_lines.append(line)
                    continue

                out_lines.append(line)
                if line.startswith("[CODEC-RESULT-URI]"):
                    payload = unquote(line[len("[CODEC-RESULT-URI]"):])
                    result = json.loads(payload)
                    self._terminate_proc(proc)
                    return result
                if line.startswith("[CODEC-RESULT]"):
                    result = json.loads(line[len("[CODEC-RESULT]"):])
                    self._terminate_proc(proc)
                    return result
                if line.startswith("[CODEC-ERROR]"):
                    self._terminate_proc(proc)
                    raise RuntimeError(line)

            # 到点未拿到结果：强杀子进程，抛 TimeoutError 让上层重试
            self._terminate_proc(proc)
            # 抓一下 err 里剩下的（可能已经通过 queue 收了大部分，这里兜底）
            drained_out, drained_err = self._drain_queue(io_q)
            out_lines.extend(drained_out)
            err_lines.extend(drained_err)
            out = "\n".join(out_lines)
            err = "\n".join(err_lines)
            raise TimeoutError(
                f"cli codec call timeout after {deadline_total}s, no result line\n"
                + "stdout:\n" + out + "\n"
                + "stderr:\n" + err
            )
        finally:
            self._terminate_proc(proc)
            try:
                os.remove(tmp_path)
            except Exception:
                pass

    @staticmethod
    def _drain_queue(io_q: "queue.Queue"):
        """把 queue 里现存的数据一次性抽干，用于超时后收拾残余输出。

        Args:
            io_q (queue.Queue): _call_via_cli 中用来收集 stdout/stderr 的队列。

        Returns:
            tuple[list[str], list[str]]: (stdout 行列表, stderr 行列表)。
        """
        outs, errs = [], []
        while True:
            try:
                tag, line = io_q.get_nowait()
            except queue.Empty:
                break
            if line is None:
                continue
            (outs if tag == "out" else errs).append(line.rstrip("\r\n"))
        return outs, errs

    @staticmethod
    def _terminate_proc(proc) -> None:
        """尽力结束 frida CLI 子进程，先 terminate 再 kill。

        Args:
            proc: `subprocess.Popen` 对象，允许为 None。
        """
        if proc is None:
            return
        if proc.poll() is not None:
            return
        try:
            proc.terminate()
        except Exception:
            pass
        try:
            proc.wait(timeout=2)
        except Exception:
            try:
                proc.kill()
            except Exception:
                pass


def encrypt_url(url: str, package: str = PKG, device_id: Optional[str] = None) -> Dict[str, Any]:
    """
    便捷函数：创建临时客户端，执行 URL 加密并自动释放资源。

    Args:
        url: 原始请求 URL。
        package: 目标应用包名。
        device_id: Frida 设备 ID（如 emulator-5554），可不传。

    Returns:
        Dict[str, Any]: 加密结果字典。
    """
    with JhsRawCodecClient(package=package, device_id=device_id) as client:
        return client.encrypt(url)


def decrypt_raw_data(
    request_url: str,
    response_raw_data: str,
    package: str = PKG,
    device_id: Optional[str] = None,
) -> Dict[str, Any]:
    """
    便捷函数：创建临时客户端，执行 raw_data 解密并自动释放资源。

    Args:
        request_url: 包含 raw_data 参数的请求 URL。
        response_raw_data: 响应中的加密 raw_data。
        package: 目标应用包名。
        device_id: Frida 设备 ID（如 emulator-5554），可不传。

    Returns:
        Dict[str, Any]: 解密结果字典。
    """
    with JhsRawCodecClient(package=package, device_id=device_id) as client:
        return client.decrypt(request_url, response_raw_data)


def call_codec(
    params: Dict[str, Any],
    package: str = PKG,
    device_id: Optional[str] = None,
) -> Dict[str, Any]:
    """
    对外统一调用入口：根据 params["op"] 执行 enc/dec。

    Args:
        params: 调用参数字典，支持两种格式：
            enc: {"op": "enc", "url": "..."}
            dec: {"op": "dec", "request_url": "...", "response_raw_data": "..."}
        package: 目标应用包名。
        device_id: Frida 设备 ID（如 emulator-5554），可不传。
            也可通过环境变量 FRIDA_DEVICE_ID 指定。

    Returns:
        Dict[str, Any]: 编解码结果字典。

    Raises:
        TypeError: 当 params 不是 dict 时抛出。
    """
    if not isinstance(params, dict):
        raise TypeError("params must be a dict")
    with JhsRawCodecClient(package=package, device_id=device_id) as client:
        return client.call(params)