auto_control/core/adb_helper.py

"""adb 命令工具封装。

并发关键：adb_connect 绝不 kill-server（会断开其他 worker 正在用的设备），
用全局锁串行化连接阶段，connect 本身很快，串行不影响整体并发。
"""
import re
import time
import subprocess
import threading

from config import ADB_PATH
from core.logger import get_logger

_log = get_logger("core.adb")

_ADB_LOCK = threading.Lock()
# adb 命令超时（秒）。adb connect 到不可达主机（如 Tailscale IP 不通）会长时间挂起，
# 没有超时会让 worker 永久卡在 connecting。connect/disconnect 正常都应瞬时完成。
_ADB_TIMEOUT = 30


def _adb(*args):
    """执行 adb 命令，返回 (stdout+stderr) 文本。不用 shell=True，避免子进程环境问题。

    用 bytes 模式自己解码（Windows 下 subprocess text=True 会用 GBK 读 adb 输出，
    遇到非 GBK 字节会崩溃）。
    带 30s 超时：adb connect 到不可达地址可能挂很久，超时返回空串由调用方处理。
    """
    try:
        r = subprocess.run([ADB_PATH, *args], capture_output=True, timeout=_ADB_TIMEOUT)
    except subprocess.TimeoutExpired:
        _log.warning(f"adb {' '.join(args)} 超时 {_ADB_TIMEOUT}s")
        return ""
    out = (r.stdout or b"") + (r.stderr or b"")
    try:
        return out.decode("utf-8", errors="replace").strip()
    except Exception:
        return out.decode("gbk", errors="replace").strip()


def get_foreground_app(remote_adb_url, timeout=3):
    """获取设备当前前台 App 包名。

    通过 adb shell dumpsys window 解析 mCurrentFocus。
    remote_adb_url: adb 连接地址（如 192.168.20.17:12345），设备需已 adb connect。
    返回包名字符串（如 com.ss.android.ugc.aweme），失败/无设备返回 None。
    """
    if not remote_adb_url:
        return None
    try:
        r = subprocess.run(
            [ADB_PATH, "-s", remote_adb_url, "shell",
             "dumpsys window | grep mCurrentFocus"],
            capture_output=True, timeout=timeout
        )
        out = (r.stdout or b"").decode("utf-8", errors="replace").strip()
        if not out:
            return None
        low = out.lower()
        if "device offline" in low or "device not found" in low or "error:" in low:
            return None
        # 解析 mCurrentFocus=Window{abc u0 com.xx/com.xx.MainActivity}
        for line in out.split("\n"):
            if "mCurrentFocus" not in line:
                continue
            # 提取 u0 后面的 token（包名/Activity 或包名）
            m = re.search(r'\bu0\s+(\S+)', line)
            if m:
                token = m.group(1).rstrip("}")
                if "/" in token:
                    return token.split("/")[0]
                return token or None
        return None
    except subprocess.TimeoutExpired:
        return None
    except Exception:
        return None


def adb_connect(remote_adb_url, retries=5):
    """adb connect（带重试）。返回 True/False。

    首次常因认证握手未就绪报 failed to authenticate，第2次重试通常 already connected。
    STF remoteConnect 启用后端口监听可能需要几秒，所以重试间隔递增。

    **绝不 kill-server**：kill-server 会断开所有设备的 adb transport，
    导致 STF provider 对全部设备误判离线并触发重连，影响所有运行中的任务。
    连接失败就返回 False，由调用方处理（标记任务失败，不影响其他设备）。
    """
    with _ADB_LOCK:
        return _try_connect(remote_adb_url, retries)


def _try_connect(remote_adb_url, retries):
    """尝试连接，不 kill-server。返回 True/False。"""
    for attempt in range(1, retries + 1):
        out = _adb("connect", remote_adb_url)
        _log.info(f"{remote_adb_url} 第 {attempt} 次尝试: {out}")
        if "connected" in out.lower() and "failed" not in out.lower():
            return True
        time.sleep(1 + attempt)
    return False


def adb_connect_light(remote_adb_url):
    """轻量 adb connect（扫描专用）：单次尝试，不重试，绝不 kill-server。

    与 adb_connect 的区别：
      - 不重试：失败立即返回 False
      - 不 kill-server：避免断开其他 worker 正在使用的 ADB 连接
    用于前台 App 扫描等非关键路径，连不上就跳过，不影响正在运行的任务。
    """
    with _ADB_LOCK:
        out = _adb("connect", remote_adb_url)
        if "connected" in out.lower() and "failed" not in out.lower():
            return True
        return False


def adb_disconnect(remote_adb_url):
    """adb disconnect。"""
    try:
        _adb("disconnect", remote_adb_url)
    except Exception:
        pass


def screenshot(serial, timeout=5):
    """获取设备当前画面截图（PNG bytes）。

    用 `adb exec-out -s <serial> shell screencap -p`，只读操作，
    不抢占 u2 的 atx-agent 通道，任务运行中调用安全。
    返回 (ok, png_bytes_or_errmsg)。
    """
    if not serial:
        return False, "缺少 serial"
    # 先尝试连一下（任务运行中设备已 connect，这里是幂等的）
    if not adb_connect_light(serial):
        return False, "adb connect 失败"
    try:
        r = subprocess.run(
            [ADB_PATH, "-s", serial, "exec-out", "screencap", "-p"],
            capture_output=True, timeout=timeout,
        )
        # PNG 文件头：\x89PNG
        if r.stdout and r.stdout[:4] == b"\x89PNG":
            return True, r.stdout
        # 失败时 stdout 可能为空，看 stderr
        err = (r.stderr or b"").decode("utf-8", errors="replace").strip()
        return False, err or "截图失败（设备可能离线）"
    except subprocess.TimeoutExpired:
        return False, "截图超时"
    except Exception as e:
        return False, str(e)


def identify_device(serial, duration=3):
    """让设备响铃以供识别（直连 serial，不经过 STF）。

    播放铃声 duration 秒后自动停止。
    不 disconnect（避免断开 STF agent 共享的 adb 连接）。
    返回 (ok, msg)。
    """
    import subprocess
    if not adb_connect_light(serial):
        return False, "adb connect 失败"
    try:
        # 异步启动铃声播放（cmd audio play-ringtone 会阻塞，用 Popen）
        p = subprocess.Popen(
            [ADB_PATH, "-s", serial, "shell", "cmd", "audio", "play-ringtone"],
            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
        )
        time.sleep(duration)
        p.terminate()
        # 停止铃声
        subprocess.run(
            [ADB_PATH, "-s", serial, "shell", "cmd", "audio", "stop-ringtone"],
            capture_output=True, timeout=5,
        )
        _log.info(f"[{serial}] 识别响铃 {duration}s 完成")
        return True, f"已响铃 {duration} 秒"
    except Exception as e:
        _log.warning(f"[{serial}] 识别响铃失败: {e}")
        return False, str(e)