349 lines
15 KiB
Python
349 lines
15 KiB
Python
"""Task-scoped Android files and native-resolution screenshot artifacts."""
|
||
|
||
import asyncio
|
||
import ctypes
|
||
from datetime import datetime, UTC
|
||
import hashlib
|
||
from io import BytesIO
|
||
import json
|
||
import os
|
||
from pathlib import Path, PurePosixPath
|
||
import re
|
||
import shlex
|
||
from typing import Any, Literal
|
||
import uuid
|
||
|
||
from PIL import Image
|
||
from pydantic import BaseModel, Field
|
||
|
||
from artemis.core.tool_failure import ToolFailure
|
||
from artemis.runtime.adb_endpoint import AdbSession, current_adb_endpoint
|
||
from artemis.tools.base import ArtemisTool
|
||
|
||
|
||
# Bound memory, model context and subprocess lifetime independently of task timeout.
|
||
MAX_FILE_BYTES = 64 * 1024 * 1024
|
||
MAX_TEXT_BYTES = 64 * 1024
|
||
ADB_TIMEOUT_SECONDS = 60
|
||
|
||
|
||
def android_files_available(ctx: Any) -> bool:
|
||
engine = getattr(ctx, "data_engine", None)
|
||
return bool(
|
||
getattr(ctx, "execution_target", "android") != "computer"
|
||
and getattr(getattr(ctx, "device", None), "device_id", None)
|
||
and engine is not None
|
||
and getattr(engine, "current_session_id", None)
|
||
)
|
||
|
||
|
||
def desktop_directory() -> Path:
|
||
"""Resolve the actual Windows desktop, including redirected/OneDrive desktops."""
|
||
if os.name == "nt":
|
||
buffer = ctypes.create_unicode_buffer(32768)
|
||
if ctypes.windll.shell32.SHGetFolderPathW(None, 0x10, None, 0, buffer) != 0:
|
||
raise OSError("无法定位当前用户的桌面目录")
|
||
path = Path(buffer.value)
|
||
else:
|
||
path = Path.home() / "Desktop"
|
||
if not path.is_dir():
|
||
raise OSError("桌面目录不存在；请使用 destination=trace")
|
||
return path.resolve()
|
||
|
||
|
||
def _filename(value: str) -> str:
|
||
if (
|
||
not value
|
||
or len(value) > 180
|
||
or value in {".", ".."}
|
||
or value != value.strip(" .")
|
||
or re.search(r'[\\/:*?"<>|\x00-\x1f]', value)
|
||
or re.match(r"^(CON|PRN|AUX|NUL|COM[0-9]|LPT[0-9])(?:\.|$)", value, re.I)
|
||
or value.startswith(".artemis-")
|
||
):
|
||
raise ValueError("filename 必须是普通文件名，不能包含路径或 Windows 保留名称")
|
||
return value
|
||
|
||
|
||
def _device_path(value: str) -> str:
|
||
path = PurePosixPath(value)
|
||
if (
|
||
not value.startswith("/")
|
||
or value.startswith("//")
|
||
or path == PurePosixPath("/")
|
||
or ".." in path.parts
|
||
or any(ord(char) < 32 for char in value)
|
||
):
|
||
raise ValueError("device_path 必须是 Android 绝对路径，不能包含 .. 或控制字符")
|
||
return str(path)
|
||
|
||
|
||
def _host_path(root: Path, name: str) -> Path:
|
||
path = root / _filename(name)
|
||
if path.is_symlink() or path.resolve().parent != root.resolve():
|
||
raise ValueError("文件不能是符号链接或指向目标目录之外")
|
||
return path
|
||
|
||
|
||
class AndroidFiles:
|
||
"""Use only the task's selected serial and its process-pinned ADB endpoint."""
|
||
|
||
def __init__(self, ctx: Any):
|
||
if not android_files_available(ctx):
|
||
raise ValueError("安卓文件工具需要当前任务的 Android 设备和 trace 会话")
|
||
self.serial = str(ctx.device.device_id)
|
||
self.root = Path(ctx.data_engine.base_dir) / "artifacts" / "android"
|
||
self.adb = AdbSession(current_adb_endpoint())
|
||
|
||
async def _run(self, *args: str, limit: int = MAX_TEXT_BYTES) -> bytes:
|
||
process = await asyncio.create_subprocess_exec(
|
||
*self.adb.command(["-s", self.serial, *args]),
|
||
env=self.adb.environment(),
|
||
stdin=asyncio.subprocess.DEVNULL,
|
||
stdout=asyncio.subprocess.PIPE,
|
||
stderr=asyncio.subprocess.PIPE,
|
||
**({"creationflags": 0x08000000} if os.name == "nt" else {}),
|
||
)
|
||
|
||
async def bounded_read(stream, maximum):
|
||
output = bytearray()
|
||
while chunk := await stream.read(65536):
|
||
output.extend(chunk)
|
||
if len(output) > maximum:
|
||
raise ValueError(f"ADB 输出超过 {maximum} 字节限制")
|
||
return bytes(output)
|
||
|
||
readers = [
|
||
asyncio.create_task(bounded_read(process.stdout, limit)),
|
||
asyncio.create_task(bounded_read(process.stderr, MAX_TEXT_BYTES)),
|
||
]
|
||
try:
|
||
async with asyncio.timeout(ADB_TIMEOUT_SECONDS):
|
||
stdout, stderr = await asyncio.gather(*readers)
|
||
await process.wait()
|
||
if process.returncode:
|
||
detail = (stderr or stdout).decode("utf-8", errors="replace")[:2000]
|
||
raise OSError(f"ADB 文件操作失败（{process.returncode}）：{detail}")
|
||
return stdout
|
||
finally:
|
||
if process.returncode is None:
|
||
try:
|
||
process.kill()
|
||
except ProcessLookupError:
|
||
pass
|
||
await process.wait()
|
||
for reader in readers:
|
||
if not reader.done():
|
||
reader.cancel()
|
||
await asyncio.gather(*readers, return_exceptions=True)
|
||
|
||
async def _shell(self, command: str) -> str:
|
||
return (await self._run("shell", command)).decode("utf-8", errors="strict").strip()
|
||
|
||
async def _hash(self, path: str) -> str:
|
||
result = await self._shell(f"sha256sum {shlex.quote(path)}")
|
||
digest = result.split(maxsplit=1)[0] if result else ""
|
||
if not re.fullmatch(r"[0-9a-fA-F]{64}", digest):
|
||
raise OSError("设备未返回有效 SHA-256，无法确认文件内容")
|
||
return digest.lower()
|
||
|
||
async def _read(self, path: str, limit: int) -> bytes:
|
||
quoted = shlex.quote(path)
|
||
size = int(await self._shell(f"test -f {quoted} && stat -c %s {quoted}"))
|
||
if not 0 <= size <= limit:
|
||
raise ValueError(f"文件大小 {size} 超过上限 {limit} 字节，请缩小文件或使用 pull")
|
||
before = await self._hash(path)
|
||
data = await self._run("exec-out", f"cat {quoted}", limit=limit)
|
||
if len(data) != size or hashlib.sha256(data).hexdigest() != before:
|
||
raise OSError("文件传输不完整或传输期间内容发生变化，未保存")
|
||
if await self._hash(path) != before:
|
||
raise OSError("传输期间设备文件发生变化，未保存")
|
||
return data
|
||
|
||
def _save(self, data: bytes, name: str, destination: str, overwrite: bool) -> dict:
|
||
self.root.mkdir(parents=True, exist_ok=True)
|
||
trace_path = _host_path(self.root, name)
|
||
targets = [trace_path]
|
||
if destination == "desktop":
|
||
targets.append(_host_path(desktop_directory(), name))
|
||
for path in targets:
|
||
if path.exists() and not overwrite:
|
||
raise FileExistsError(f"文件已存在：{path}；换名或在获准覆盖后使用 overwrite=true")
|
||
digest = hashlib.sha256(data).hexdigest()
|
||
for path in targets:
|
||
temporary = path.parent / f".artemis-{uuid.uuid4().hex}.tmp"
|
||
try:
|
||
with temporary.open("xb") as handle:
|
||
handle.write(data)
|
||
handle.flush()
|
||
os.fsync(handle.fileno())
|
||
if hashlib.sha256(temporary.read_bytes()).hexdigest() != digest:
|
||
raise OSError(f"保存后读回校验失败：{path}")
|
||
if overwrite:
|
||
os.replace(temporary, path)
|
||
elif os.name == "nt":
|
||
# Windows rename is exclusive and also works on exFAT/FAT.
|
||
os.rename(temporary, path)
|
||
else:
|
||
# Publish only a complete file, without clobbering a racing writer.
|
||
os.link(temporary, path)
|
||
finally:
|
||
temporary.unlink(missing_ok=True)
|
||
return {
|
||
"path": str(targets[-1]),
|
||
"trace_path": str(trace_path),
|
||
"bytes": len(data),
|
||
"sha256": digest,
|
||
"verified": True,
|
||
}
|
||
|
||
async def _write(self, path: str, data: bytes, overwrite: bool) -> dict:
|
||
quoted = shlex.quote(path)
|
||
parent = str(PurePosixPath(path).parent)
|
||
remote_temp = f"{parent}/.artemis-{uuid.uuid4().hex}.tmp"
|
||
temp_quoted = shlex.quote(remote_temp)
|
||
self.root.mkdir(parents=True, exist_ok=True)
|
||
local_temp = self.root / f".artemis-{uuid.uuid4().hex}.tmp"
|
||
if not overwrite:
|
||
exists = await self._shell(
|
||
f"if test -e {quoted} || test -L {quoted}; then printf exists; fi"
|
||
)
|
||
if exists:
|
||
raise FileExistsError(f"设备文件已存在：{path}；换名或获准后使用 overwrite=true")
|
||
await self._shell(f"mkdir -p {shlex.quote(parent)}")
|
||
digest = hashlib.sha256(data).hexdigest()
|
||
local_temp.write_bytes(data)
|
||
try:
|
||
await self._run("push", str(local_temp), remote_temp)
|
||
if await self._hash(remote_temp) != digest:
|
||
raise OSError("写入设备临时文件后校验失败，未替换目标文件")
|
||
# Toybox mv -n returns success when it skipped an existing target;
|
||
# the temporary file's continued existence makes that a failure.
|
||
flag = "-f" if overwrite else "-n"
|
||
await self._shell(
|
||
f"test ! -d {quoted} && test ! -L {quoted} && "
|
||
f"mv {flag} {temp_quoted} {quoted} && test ! -e {temp_quoted}"
|
||
)
|
||
if await self._hash(path) != digest:
|
||
raise OSError("设备目标文件读回校验失败")
|
||
return {"device_path": path, "bytes": len(data), "sha256": digest, "verified": True}
|
||
finally:
|
||
local_temp.unlink(missing_ok=True)
|
||
# Remove only this operation's own temporary file, even on failure.
|
||
try:
|
||
await self._shell(f"rm -f {temp_quoted}")
|
||
except (OSError, ValueError, TimeoutError):
|
||
pass
|
||
|
||
async def execute(self, args: "AndroidFileArgs") -> dict:
|
||
if args.action == "screenshot":
|
||
name = _filename(args.filename or f"android-{uuid.uuid4().hex[:12]}.png")
|
||
if Path(name).suffix.lower() != ".png":
|
||
raise ValueError("原生截图保存为 PNG，filename 必须以 .png 结尾")
|
||
data = await self._run("exec-out", "screencap", "-p", limit=MAX_FILE_BYTES)
|
||
with Image.open(BytesIO(data)) as image:
|
||
if image.format != "PNG":
|
||
raise ValueError("设备没有返回有效 PNG 截图")
|
||
image.load()
|
||
dimensions = {"width": image.width, "height": image.height}
|
||
return {
|
||
**self._save(data, name, args.destination, args.overwrite),
|
||
**dimensions,
|
||
"captured_at": datetime.now(UTC).isoformat(),
|
||
"device_serial": self.serial,
|
||
}
|
||
path = _device_path(args.device_path)
|
||
if args.action == "list":
|
||
return {
|
||
"device_path": path,
|
||
"entries": await self._shell(f"ls -la {shlex.quote(path)}"),
|
||
}
|
||
if args.action == "read_text":
|
||
data = await self._read(path, MAX_TEXT_BYTES)
|
||
return {"device_path": path, "content": data.decode("utf-8-sig"), "bytes": len(data)}
|
||
if args.action == "pull":
|
||
data = await self._read(path, MAX_FILE_BYTES)
|
||
return self._save(
|
||
data, args.filename or PurePosixPath(path).name, args.destination, args.overwrite
|
||
)
|
||
if args.action == "write_text":
|
||
if args.content is None:
|
||
raise ValueError("write_text 需要 content（UTF-8 文本，可为空字符串）")
|
||
data = args.content.encode("utf-8")
|
||
if len(data) > MAX_TEXT_BYTES:
|
||
raise ValueError("文本超过 64 KiB，请使用 push 传输现有文件")
|
||
else: # push: only a named artifact or a file explicitly selected on the desktop
|
||
root = self.root if args.destination == "trace" else desktop_directory()
|
||
source = _host_path(root, args.filename)
|
||
if not source.is_file() or source.stat().st_size > MAX_FILE_BYTES:
|
||
raise ValueError("本机源文件不存在或超过 64 MiB")
|
||
with source.open("rb") as handle:
|
||
data = handle.read(MAX_FILE_BYTES + 1)
|
||
if len(data) > MAX_FILE_BYTES:
|
||
raise ValueError("本机源文件超过 64 MiB")
|
||
return await self._write(path, data, args.overwrite)
|
||
|
||
|
||
class AndroidFileArgs(BaseModel):
|
||
action: Literal["screenshot", "list", "read_text", "write_text", "pull", "push"] = Field(
|
||
description="screenshot captures Android to host; list/read_text/write_text use Android paths; pull downloads, push uploads."
|
||
)
|
||
device_path: str = Field(
|
||
default="",
|
||
description="Absolute Android path, e.g. /sdcard/Download/result.txt. Not a PC path. Unused for screenshot.",
|
||
)
|
||
filename: str = Field(
|
||
default="",
|
||
description="Host file name only, no directories. Screenshot: .png (generated if empty); pull: defaults to device basename; push: required existing file.",
|
||
)
|
||
destination: Literal["trace", "desktop"] = Field(
|
||
default="trace",
|
||
description="Host location. Default: current task trace/artifacts/android. desktop: current user's actual desktop, only if requested; downloads also keep a trace copy. For push this selects the source location.",
|
||
)
|
||
content: str | None = Field(
|
||
default=None,
|
||
description="UTF-8 text for write_text only, max 64 KiB; empty string creates an empty file.",
|
||
)
|
||
overwrite: bool = Field(
|
||
default=False,
|
||
description="Only true when the user authorized replacing an existing file; default never overwrites.",
|
||
)
|
||
|
||
|
||
async def _android_file(ctx: Any, **kwargs: Any) -> str:
|
||
try:
|
||
args = AndroidFileArgs.model_validate(kwargs)
|
||
result = await AndroidFiles(ctx).execute(args)
|
||
return json.dumps({"ok": True, "action": args.action, **result}, ensure_ascii=False)
|
||
except TimeoutError:
|
||
return ToolFailure(f"安卓文件操作超时（{ADB_TIMEOUT_SECONDS} 秒），ADB 子进程已停止")
|
||
except (OSError, ValueError) as exc:
|
||
return ToolFailure(f"安卓文件操作失败：{exc or type(exc).__name__}")
|
||
|
||
|
||
class AndroidFileTool(ArtemisTool):
|
||
def is_available(self, ctx: Any = None) -> bool:
|
||
return android_files_available(ctx)
|
||
|
||
async def execute(self, driver=None, ctx=None, **kwargs):
|
||
return await _android_file(ctx, **kwargs)
|
||
|
||
|
||
android_file = AndroidFileTool(
|
||
name="android_file",
|
||
description=(
|
||
"Save Android screenshots and read/write/transfer Android files on the task's selected device. "
|
||
"Use action=screenshot to save the CURRENT phone screen as native PNG, without UI shortcuts "
|
||
"or the computer-use plugin. Default host location is this task's trace/artifacts/android; "
|
||
"use destination=desktop only when the user requests the PC desktop. Returns verified paths, "
|
||
"size and SHA-256; never claim success without a successful result. Screenshot capture does "
|
||
"not prove the page is correct. Capture only after intended UI actions have executed and "
|
||
"the resulting screen has been observed, not alongside pending UI actions. "
|
||
"list/read_text/write_text operate on absolute ANDROID paths; "
|
||
"text is UTF-8, max 64 KiB. pull/push transfer a single file up to 64 MiB. No root or app-private "
|
||
"storage bypass; report Android permission errors. Existing files are protected by default."
|
||
),
|
||
args_schema=AndroidFileArgs,
|
||
category="system",
|
||
)
|