286 lines
11 KiB
Python
286 lines
11 KiB
Python
"""Bounded TCP device discovery shared by the console, CLI and MCP.
|
|
|
|
USB and Android SDK emulators are discovered by ADB itself. Third-party
|
|
emulators expose TCP adbd listeners; only listeners owned by known emulator
|
|
processes are probed. Explicitly connected TCP devices are remembered separately.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import socket
|
|
import struct
|
|
import subprocess
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
|
|
import psutil
|
|
|
|
from artemis.config import settings
|
|
from artemis.config.paths import get_app_dir
|
|
from artemis.runtime.adb_endpoint import AdbSession, current_adb_endpoint
|
|
from artemis.utils.logger import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
def _canonical_target(target: str) -> str:
|
|
host, separator, port = target.rpartition(":")
|
|
if separator and host.lower() in {"localhost", "127.0.0.1"}:
|
|
return f"127.0.0.1:{port}"
|
|
return target
|
|
|
|
|
|
def emulator_listeners() -> list[list[str]]:
|
|
"""Return listener groups, one per running emulator (never scan a network)."""
|
|
owners = set()
|
|
for process in psutil.process_iter(["pid", "name", "exe"]):
|
|
info = process.info
|
|
name = str(info.get("name") or "").lower().removesuffix(".exe")
|
|
path = str(info.get("exe") or "").lower()
|
|
known = name in {
|
|
"mumunxdevice",
|
|
"nemuheadless",
|
|
"mumuheadless",
|
|
"ldvboxheadless",
|
|
"noxvmhandle",
|
|
"memuheadless",
|
|
"hd-player",
|
|
}
|
|
vendor_vm = name == "vboxheadless" and any(
|
|
vendor in path for vendor in ("mumu", "nemu", "ldplayer", "leidian", "nox", "memu")
|
|
)
|
|
if known or vendor_vm:
|
|
owners.add(info["pid"])
|
|
if not owners:
|
|
return []
|
|
groups: dict[int, set[int]] = {}
|
|
for connection in psutil.net_connections(kind="tcp"):
|
|
address = connection.laddr
|
|
pid = connection.pid
|
|
if (
|
|
pid is not None
|
|
and pid in owners
|
|
and address
|
|
and connection.status == psutil.CONN_LISTEN
|
|
and address.ip in {"127.0.0.1", "0.0.0.0"}
|
|
):
|
|
groups.setdefault(pid, set()).add(address.port)
|
|
# Prefer current MuMu ports over its compatibility aliases (same adbd).
|
|
return [
|
|
[f"127.0.0.1:{port}" for port in sorted(ports, key=lambda p: (p in {5555, 7555}, p))]
|
|
for ports in groups.values()
|
|
]
|
|
|
|
|
|
def is_adbd(target: str, timeout: float) -> bool:
|
|
"""Confirm an ADB CNXN/AUTH header before giving a local port to adb connect."""
|
|
host, port = target.rsplit(":", 1)
|
|
payload = b"host::\x00"
|
|
command = int.from_bytes(b"CNXN", "little")
|
|
packet = struct.pack(
|
|
"<6I", command, 0x01000000, 4096, len(payload), sum(payload), command ^ 0xFFFFFFFF
|
|
)
|
|
try:
|
|
with socket.create_connection((host, int(port)), timeout=timeout) as stream:
|
|
stream.sendall(packet + payload)
|
|
header = b""
|
|
deadline = time.monotonic() + timeout
|
|
while len(header) < 24:
|
|
remaining = deadline - time.monotonic()
|
|
if remaining <= 0:
|
|
return False
|
|
stream.settimeout(remaining)
|
|
chunk = stream.recv(24 - len(header))
|
|
if not chunk:
|
|
return False
|
|
header += chunk
|
|
kind, _, _, length, _, magic = struct.unpack("<6I", header)
|
|
return (
|
|
kind.to_bytes(4, "little") in {b"CNXN", b"AUTH"}
|
|
and magic == kind ^ 0xFFFFFFFF
|
|
and length <= 1048576
|
|
)
|
|
except (OSError, ValueError):
|
|
return False
|
|
|
|
|
|
class AdbDeviceDiscovery:
|
|
"""Coalesce discovery and reconnect only missing transports, with a total budget."""
|
|
|
|
def __init__(self):
|
|
self._gate = threading.Lock()
|
|
self._last_scan = float("-inf")
|
|
self.results: list[dict[str, object]] = []
|
|
self._cursor = 0
|
|
|
|
@staticmethod
|
|
def _directory() -> Path:
|
|
return get_app_dir() / "adb-devices"
|
|
|
|
def remember(self, target: str) -> bool:
|
|
"""Save an explicitly connected device, without credentials or shared-list races."""
|
|
directory = self._directory()
|
|
temporary = None
|
|
try:
|
|
directory.mkdir(parents=True, exist_ok=True)
|
|
destination = directory / (hashlib.sha256(target.encode()).hexdigest() + ".json")
|
|
fd, temporary = tempfile.mkstemp(dir=directory, prefix=".device-")
|
|
with os.fdopen(fd, "w", encoding="utf-8") as stream:
|
|
json.dump({"target": target}, stream)
|
|
os.replace(temporary, destination)
|
|
return True
|
|
except OSError as exc:
|
|
logger.warning(f"Cannot remember ADB device {target}: {exc}")
|
|
return False
|
|
finally:
|
|
if temporary:
|
|
Path(temporary).unlink(missing_ok=True)
|
|
|
|
def _remembered(self) -> list[str]:
|
|
targets = []
|
|
for path in self._directory().glob("*.json"):
|
|
try:
|
|
target = json.loads(path.read_text(encoding="utf-8"))["target"]
|
|
host, port = target.rsplit(":", 1)
|
|
if host and 0 < int(port) < 65536:
|
|
targets.append(target)
|
|
except (OSError, ValueError, KeyError, TypeError, AttributeError) as exc:
|
|
logger.warning(f"Invalid remembered ADB device {path.name}: {exc}")
|
|
return targets
|
|
|
|
def refresh(self, adb_path: str, *, force: bool = False) -> list[dict[str, object]]:
|
|
if not settings.ARTEMIS_ADB_AUTO_CONNECT or not current_adb_endpoint().is_local_default:
|
|
self.results = []
|
|
return []
|
|
if not self._gate.acquire(blocking=False):
|
|
return self.results.copy()
|
|
try:
|
|
if (
|
|
not force
|
|
and time.monotonic() - self._last_scan
|
|
< settings.ARTEMIS_ADB_DISCOVERY_INTERVAL_SECONDS
|
|
):
|
|
return self.results.copy()
|
|
self._last_scan = time.monotonic()
|
|
self.results = []
|
|
self._refresh(adb_path)
|
|
except (OSError, psutil.Error, ValueError, subprocess.SubprocessError) as exc:
|
|
logger.warning(f"ADB automatic discovery unavailable: {exc}")
|
|
self.results.append({"success": False, "message": str(exc)})
|
|
finally:
|
|
self._gate.release()
|
|
return self.results.copy()
|
|
|
|
def _refresh(self, adb_path: str) -> None:
|
|
session = AdbSession(current_adb_endpoint(), adb_path=adb_path)
|
|
deadline = time.monotonic() + settings.ARTEMIS_ADB_DISCOVERY_TIMEOUT_SECONDS
|
|
devices = session.run(
|
|
["devices"],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=max(deadline - time.monotonic(), 0.01),
|
|
)
|
|
if devices.returncode:
|
|
return
|
|
ready = {
|
|
_canonical_target(parts[0])
|
|
for line in devices.stdout.splitlines()
|
|
if len(parts := line.split()) >= 2 and parts[1] in {"device", "unauthorized"}
|
|
}
|
|
try:
|
|
groups = emulator_listeners()
|
|
except (OSError, psutil.Error) as exc:
|
|
# Missing process-inspection permissions must not stop reconnection
|
|
# of explicitly configured phones or make device enumeration fail.
|
|
groups = []
|
|
logger.warning(f"Cannot inspect local emulator listeners: {exc}")
|
|
self.results.append({"success": False, "message": str(exc)})
|
|
remembered = self._remembered()
|
|
targets = [(group, True) for group in groups] + [([target], False) for target in remembered]
|
|
# Resume at a different group if a slow/offline peer exhausts the budget.
|
|
# Large multi-instance installations must not starve later devices.
|
|
count = len(targets)
|
|
start = self._cursor % count if count else 0
|
|
for offset in range(count):
|
|
index = (start + offset) % count
|
|
candidates, probe = targets[index]
|
|
self._cursor = (index + 1) % count
|
|
if any(_canonical_target(target) in ready for target in candidates):
|
|
continue
|
|
for target in candidates:
|
|
remaining = deadline - time.monotonic()
|
|
if remaining <= 0:
|
|
return
|
|
if probe and not is_adbd(target, min(remaining, 0.3)):
|
|
continue
|
|
remaining = deadline - time.monotonic()
|
|
if remaining <= 0:
|
|
return
|
|
try:
|
|
result = session.run(
|
|
["connect", target],
|
|
capture_output=True,
|
|
text=True,
|
|
timeout=min(remaining, 2),
|
|
)
|
|
except subprocess.TimeoutExpired:
|
|
self.results.append(
|
|
{"target": target, "success": False, "message": "ADB 连接超时"}
|
|
)
|
|
continue
|
|
output = (result.stdout + result.stderr).strip()
|
|
success = (
|
|
result.returncode == 0
|
|
and "connected to" in output.lower()
|
|
and "failed" not in output.lower()
|
|
)
|
|
state = self._wait_for_transport(session, target, deadline) if success else None
|
|
pending = success and state is None
|
|
if success and state is None:
|
|
success = False
|
|
output = f"{output}；已发起连接，等待设备就绪，下次刷新会重试"
|
|
self.results.append(
|
|
{
|
|
"target": target,
|
|
"success": success,
|
|
"state": state,
|
|
"pending": pending,
|
|
"message": output,
|
|
}
|
|
)
|
|
if success:
|
|
logger.info(f"ADB 自动连接：{target} — {output}")
|
|
ready.add(_canonical_target(target))
|
|
break
|
|
if pending:
|
|
logger.info(f"ADB 设备连接中：{target} — {output}")
|
|
break # Do not create a second alias while the first transport is settling.
|
|
logger.warning(f"ADB 自动连接失败：{target} — {output}")
|
|
|
|
@staticmethod
|
|
def _wait_for_transport(session: AdbSession, target: str, deadline: float) -> str | None:
|
|
"""adb connect can return before its transport appears in adb devices."""
|
|
while (remaining := deadline - time.monotonic()) > 0:
|
|
try:
|
|
result = session.run(["devices"], capture_output=True, text=True, timeout=remaining)
|
|
except subprocess.TimeoutExpired:
|
|
return None
|
|
for line in result.stdout.splitlines():
|
|
parts = line.split()
|
|
if (
|
|
len(parts) >= 2
|
|
and _canonical_target(parts[0]) == _canonical_target(target)
|
|
and parts[1] in {"device", "unauthorized"}
|
|
):
|
|
return parts[1]
|
|
time.sleep(min(0.1, max(deadline - time.monotonic(), 0)))
|
|
return None
|
|
|
|
|
|
adb_device_discovery = AdbDeviceDiscovery()
|