fix(device-test): contain tool processes and add reviewed preview
This commit is contained in:
chenmingxuan committed
1 parent
c98d7527a0
commit
7d31e09cb7
59 files changed
+23722
No files matched your search
@@ -0,0 +1,484 @@
|
||||
"""Bounded script execution and task-owned evidence for the device-test plugin."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from contextlib import suppress
|
||||
import hashlib
|
||||
import ipaddress
|
||||
import json
|
||||
import mimetypes
|
||||
import os
|
||||
from pathlib import Path
|
||||
import re
|
||||
import shutil
|
||||
import signal
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from typing import Any
|
||||
from urllib.parse import parse_qsl, unquote, urlsplit
|
||||
import uuid
|
||||
|
||||
from mcp.types import CallToolResult, TextContent
|
||||
|
||||
|
||||
PACKAGE = Path(__file__).resolve().parent
|
||||
SCRIPTS = PACKAGE / "skill" / "scripts"
|
||||
MAX_FILE_BYTES = 64 * 1024 * 1024
|
||||
MAX_OUTPUT_BYTES = 4 * 1024 * 1024
|
||||
SECRET_KEY = re.compile(
|
||||
r"(?:^|_)(?:password|passwd|pwd|secret|token|access_token|api_key|apikey|authorization)$",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
|
||||
class InputError(ValueError):
|
||||
"""Only constant, public messages belong in this exception."""
|
||||
|
||||
|
||||
def local_file(value: str) -> Path:
|
||||
if (
|
||||
not value
|
||||
or "://" in value
|
||||
or value.startswith(("\\\\", "//"))
|
||||
or re.search(r"[\x00-\x1f]", value)
|
||||
or not Path(value).is_absolute()
|
||||
):
|
||||
raise InputError("需要本机文件的绝对路径；不接受 URL 或网络共享。")
|
||||
path = Path(value).resolve()
|
||||
if not path.is_file():
|
||||
raise InputError("本机输入文件不存在。")
|
||||
return path
|
||||
|
||||
|
||||
def unicast_host(value: str) -> str:
|
||||
try:
|
||||
if "%" in value:
|
||||
raise ValueError
|
||||
address = ipaddress.ip_address(value)
|
||||
effective = address.ipv4_mapped if isinstance(address, ipaddress.IPv6Address) else None
|
||||
effective = effective or address
|
||||
if (
|
||||
effective.is_unspecified
|
||||
or effective.is_multicast
|
||||
or str(effective) == "255.255.255.255"
|
||||
):
|
||||
raise ValueError
|
||||
except ValueError:
|
||||
raise InputError("需要一个明确的单播 IP；不接受主机名、网段或广播地址。") from None
|
||||
return str(address)
|
||||
|
||||
|
||||
def operation_budget() -> int:
|
||||
value = os.environ.get("ARTEMIS_DEVICE_TEST_TIMEOUT", "45")
|
||||
if not re.fullmatch(r"[0-9]+", value) or not 10 <= int(value) <= 3300:
|
||||
raise InputError("插件操作时限必须是 10..3300 秒的整数。")
|
||||
return int(value)
|
||||
|
||||
|
||||
def check_budget(required: float) -> None:
|
||||
if required > operation_budget():
|
||||
raise InputError("所需时限超过插件操作预算；请缩短操作或同时调整宿主与插件时限。")
|
||||
|
||||
|
||||
def artifact_root() -> Path:
|
||||
value = os.environ.get("ARTEMIS_ARTIFACTS_DIR", "")
|
||||
if not value:
|
||||
raise InputError("缺少 Artemis 注入的任务资料目录；请通过 Artemis 任务调用。")
|
||||
path = Path(value).absolute()
|
||||
if path.resolve() != path or path.is_symlink():
|
||||
raise InputError("任务资料目录不能通过链接指向其他位置。")
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
return path
|
||||
|
||||
|
||||
def new_artifact(prefix: str, suffix: str) -> Path:
|
||||
return artifact_root() / f"{prefix}-{uuid.uuid4().hex}{suffix}"
|
||||
|
||||
|
||||
def task_evidence(name: str, *, allow_missing: bool = False) -> Path:
|
||||
if (
|
||||
not name
|
||||
or len(name) > 180
|
||||
or name != Path(name).name
|
||||
or any(char in name for char in '/\\:*?"<>|')
|
||||
or name in {".", ".."}
|
||||
or name.rstrip(" .") != name
|
||||
or re.search(r"[\x00-\x1f]", name)
|
||||
):
|
||||
raise InputError("证据必须引用当前任务中的普通附件文件名。")
|
||||
path = artifact_root() / name
|
||||
if path.is_symlink() or path.resolve().parent != artifact_root():
|
||||
raise InputError("证据文件不存在或超出当前任务资料目录。")
|
||||
if not path.is_file() and not (allow_missing and not path.exists()):
|
||||
raise InputError("证据文件不存在或超出当前任务资料目录。")
|
||||
if path.is_file() and path.stat().st_size > MAX_FILE_BYTES:
|
||||
raise InputError("单份证据不能超过 64 MiB。")
|
||||
return path
|
||||
|
||||
|
||||
def sha256(path: Path) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as stream:
|
||||
for block in iter(lambda: stream.read(1024 * 1024), b""):
|
||||
digest.update(block)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def sanitized(value: Any, secrets: tuple[str, ...] = ()) -> Any:
|
||||
if isinstance(value, dict):
|
||||
return {
|
||||
key: "[REDACTED]" if SECRET_KEY.search(str(key)) else sanitized(item, secrets)
|
||||
for key, item in value.items()
|
||||
}
|
||||
if isinstance(value, list):
|
||||
return [sanitized(item, secrets) for item in value]
|
||||
if not isinstance(value, str):
|
||||
return value
|
||||
for secret in sorted(filter(None, secrets), key=len, reverse=True):
|
||||
if value == secret:
|
||||
return "[REDACTED]"
|
||||
if len(secret) >= 4:
|
||||
value = value.replace(secret, "[REDACTED]")
|
||||
value = re.sub(r"(?i)rtsps?://[^\s\"<>]+", "[RTSP_URL_REDACTED]", value)
|
||||
return re.sub(
|
||||
r"(?i)(password|passwd|pwd|token|secret|api_key|authorization)\s*[:=]\s*[^\s,;]+",
|
||||
r"\1=[REDACTED]",
|
||||
value,
|
||||
)
|
||||
|
||||
|
||||
def rtsp_environment(name: str) -> tuple[dict[str, str], tuple[str, ...]]:
|
||||
if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]{0,99}", name):
|
||||
raise InputError("RTSP 输入必须是环境变量名称。")
|
||||
value = os.environ.get(name, "")
|
||||
config = os.environ.get("ARTEMIS_DEVICE_TEST_ENV_FILE", "")
|
||||
if not value and config:
|
||||
from dotenv import dotenv_values
|
||||
|
||||
value = dotenv_values(local_file(config), interpolate=False).get(name) or ""
|
||||
try:
|
||||
parsed = urlsplit(value)
|
||||
if (
|
||||
parsed.scheme not in {"rtsp", "rtsps"}
|
||||
or not parsed.hostname
|
||||
or parsed.port == 0
|
||||
or parsed.fragment
|
||||
or re.search(r"[\x00-\x20\x7f]", value)
|
||||
):
|
||||
raise ValueError
|
||||
except ValueError:
|
||||
raise InputError("指定的 RTSP 环境变量未配置或格式无效；请在本机凭据环境中配置。") from None
|
||||
secrets = tuple(
|
||||
filter(
|
||||
None,
|
||||
(
|
||||
value,
|
||||
parsed.username,
|
||||
parsed.password,
|
||||
unquote(parsed.password or ""),
|
||||
*(item for _, item in parse_qsl(parsed.query)),
|
||||
),
|
||||
)
|
||||
)
|
||||
return {name: value}, secrets
|
||||
|
||||
|
||||
def runner_arguments(command: str, arguments: list[str]) -> list[str]:
|
||||
budget = operation_budget()
|
||||
result = ["--timeout", str(budget)]
|
||||
profile = os.environ.get("ARTEMIS_DEVICE_TEST_PROFILE", "")
|
||||
tool_root = os.environ.get("ARTEMIS_DEVICE_TEST_TOOL_ROOT", "")
|
||||
if profile:
|
||||
result += ["--profile", str(local_file(profile))]
|
||||
if tool_root:
|
||||
root = Path(tool_root)
|
||||
if not root.is_absolute() or not root.is_dir() or tool_root.startswith(("\\\\", "//")):
|
||||
raise InputError("工具目录必须是本机已解压目录的绝对路径。")
|
||||
result += ["--tool-root", str(root.resolve())]
|
||||
# Leave time for the runner to terminate its own child before the outer deadline.
|
||||
result[1] = str(max(1, budget - 3))
|
||||
return [*result, command, *arguments]
|
||||
|
||||
|
||||
def failure(
|
||||
message: str,
|
||||
code: str = "invalid_input",
|
||||
*,
|
||||
execution: str = "not_started",
|
||||
side_effect: str = "none",
|
||||
) -> CallToolResult:
|
||||
return CallToolResult(
|
||||
isError=True,
|
||||
content=[TextContent(type="text", text=message)],
|
||||
structuredContent={
|
||||
"test_verdict": "not_evaluated",
|
||||
"artemis": {
|
||||
"version": 1,
|
||||
"execution": execution,
|
||||
"side_effect": side_effect,
|
||||
"retryable": False,
|
||||
"error_code": code,
|
||||
"attachments": [],
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def attach_file(path: Path) -> dict:
|
||||
if path.is_symlink() or not path.is_file() or path.stat().st_size > MAX_FILE_BYTES:
|
||||
raise InputError("附件缺失、使用链接或超过 64 MiB；不能声称已归档。")
|
||||
root = artifact_root()
|
||||
if path.parent.resolve() != root:
|
||||
destination = new_artifact("evidence", path.suffix[:16])
|
||||
temporary = destination.with_name(destination.name + ".tmp")
|
||||
with path.open("rb") as source, temporary.open("xb") as target:
|
||||
shutil.copyfileobj(source, target)
|
||||
if sha256(path) != sha256(temporary):
|
||||
raise InputError("证据复制后的 SHA-256 不匹配。")
|
||||
temporary.rename(destination)
|
||||
path = destination
|
||||
return {
|
||||
"name": path.name,
|
||||
"media_type": mimetypes.guess_type(path.name)[0] or "application/octet-stream",
|
||||
"sha256": sha256(path),
|
||||
}
|
||||
|
||||
|
||||
def receipt(
|
||||
operation: str,
|
||||
data: dict,
|
||||
*,
|
||||
target: str = "local-environment",
|
||||
kind: str = "terminal",
|
||||
execution: str = "completed",
|
||||
side_effect: str = "none",
|
||||
error_code: str = "",
|
||||
files: list[Path] | None = None,
|
||||
) -> CallToolResult:
|
||||
data = sanitized(data)
|
||||
attachments = []
|
||||
issues = []
|
||||
for path in files or []:
|
||||
try:
|
||||
attachments.append(attach_file(path))
|
||||
except (InputError, OSError):
|
||||
issues.append("某份输出未能回读归档；可能缺失、使用链接或超过 64 MiB。")
|
||||
if issues:
|
||||
data["attachment_issues"] = issues
|
||||
error_code = error_code or "attachment_unavailable"
|
||||
data["test_verdict"] = "not_evaluated"
|
||||
log = new_artifact(operation, ".json")
|
||||
with log.open("x", encoding="utf-8") as stream:
|
||||
json.dump(data, stream, ensure_ascii=False, indent=2)
|
||||
attachments.append(attach_file(log))
|
||||
target_id = (
|
||||
target if len(target) <= 256 else "sha256:" + hashlib.sha256(target.encode()).hexdigest()
|
||||
)
|
||||
data["artemis"] = {
|
||||
"version": 1,
|
||||
"execution": execution,
|
||||
"side_effect": side_effect,
|
||||
"retryable": False,
|
||||
"source": {
|
||||
"kind": kind,
|
||||
"target_id": target_id,
|
||||
"label": operation,
|
||||
"captured_at": time.time(),
|
||||
},
|
||||
"attachments": attachments,
|
||||
}
|
||||
if error_code:
|
||||
data["artemis"]["error_code"] = error_code
|
||||
text = json.dumps(data, ensure_ascii=False)
|
||||
if len(text) > 16000:
|
||||
text = text[:15000] + "\n结果较长；完整结构与任务 JSON 附件保留原始字段。"
|
||||
return CallToolResult(
|
||||
isError=bool(error_code),
|
||||
content=[TextContent(type="text", text=text)],
|
||||
structuredContent=data,
|
||||
)
|
||||
|
||||
|
||||
async def stop_process(process: asyncio.subprocess.Process) -> None:
|
||||
"""Reap the worker and its fixed tool descendants before returning from cancellation."""
|
||||
if os.name == "nt":
|
||||
command = str(
|
||||
Path(os.environ.get("SystemRoot", r"C:\Windows")) / "System32" / "taskkill.exe"
|
||||
)
|
||||
killer = await asyncio.create_subprocess_exec(
|
||||
command,
|
||||
"/PID",
|
||||
str(process.pid),
|
||||
"/T",
|
||||
"/F",
|
||||
stdout=asyncio.subprocess.DEVNULL,
|
||||
stderr=asyncio.subprocess.DEVNULL,
|
||||
creationflags=subprocess.CREATE_NO_WINDOW,
|
||||
)
|
||||
try:
|
||||
await asyncio.wait_for(killer.wait(), 5)
|
||||
except TimeoutError:
|
||||
killer.kill()
|
||||
await killer.wait()
|
||||
else:
|
||||
with suppress(ProcessLookupError):
|
||||
os.killpg(process.pid, signal.SIGKILL)
|
||||
if process.returncode is None:
|
||||
with suppress(ProcessLookupError):
|
||||
process.kill()
|
||||
await process.wait()
|
||||
|
||||
|
||||
async def read_output(process: asyncio.subprocess.Process) -> bytes:
|
||||
chunks = []
|
||||
size = 0
|
||||
while chunk := await process.stdout.read(65536):
|
||||
size += len(chunk)
|
||||
if size > MAX_OUTPUT_BYTES:
|
||||
raise InputError("工具输出超过 4 MiB，已停止；部分输出不作为有效观测。")
|
||||
chunks.append(chunk)
|
||||
await process.wait()
|
||||
return b"".join(chunks)
|
||||
|
||||
|
||||
async def run_cli(
|
||||
script: str,
|
||||
arguments: list[str],
|
||||
*,
|
||||
operation: str,
|
||||
target: str = "local-environment",
|
||||
kind: str = "terminal",
|
||||
write: bool = False,
|
||||
files: list[Path] | None = None,
|
||||
env: dict[str, str] | None = None,
|
||||
secrets: tuple[str, ...] = (),
|
||||
timeout: float | None = None,
|
||||
) -> CallToolResult:
|
||||
artifact_root()
|
||||
budget = timeout or operation_budget()
|
||||
native_scripts = {"network_inventory.py", "dependency_runner.py"}
|
||||
if script not in {"tool_runner.py", "device_adapter.py", "evidence.py", *native_scripts}:
|
||||
raise InputError("不支持的脚本入口。")
|
||||
script_path = (PACKAGE if script in native_scripts else SCRIPTS) / script
|
||||
child_env = os.environ.copy()
|
||||
for name in ("FFREPORT", "PYTHONPATH", "PYTHONHOME"):
|
||||
child_env.pop(name, None)
|
||||
child_env.update(env or {})
|
||||
options = (
|
||||
{"creationflags": subprocess.CREATE_NO_WINDOW | subprocess.CREATE_NEW_PROCESS_GROUP}
|
||||
if os.name == "nt"
|
||||
else {"start_new_session": True}
|
||||
)
|
||||
try:
|
||||
process = await asyncio.create_subprocess_exec(
|
||||
sys.executable,
|
||||
"-I",
|
||||
"-B",
|
||||
str(script_path),
|
||||
*arguments,
|
||||
cwd=str(PACKAGE),
|
||||
env=child_env,
|
||||
stdin=asyncio.subprocess.DEVNULL,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.DEVNULL,
|
||||
**options,
|
||||
)
|
||||
except OSError:
|
||||
return failure("固定脚本无法启动；未执行设备动作。", "worker_launch_failed")
|
||||
try:
|
||||
output = await asyncio.wait_for(read_output(process), budget)
|
||||
except (TimeoutError, InputError) as exc:
|
||||
await stop_process(process)
|
||||
return receipt(
|
||||
operation,
|
||||
{"execution_status": "timed_out" if isinstance(exc, TimeoutError) else "output_limit"},
|
||||
target=target,
|
||||
kind=kind,
|
||||
execution="partial",
|
||||
side_effect="possible" if write else "none",
|
||||
error_code="worker_timeout" if isinstance(exc, TimeoutError) else "output_limit",
|
||||
files=files,
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
await asyncio.shield(stop_process(process))
|
||||
raise
|
||||
except OSError:
|
||||
await stop_process(process)
|
||||
return failure(
|
||||
"工具输出通道中断；执行结果待核对，不能自动重试。",
|
||||
"worker_output_interrupted",
|
||||
execution="unknown",
|
||||
side_effect="possible" if write else "none",
|
||||
)
|
||||
try:
|
||||
decoded = output.decode("utf-8")
|
||||
try:
|
||||
result = json.loads(decoded)
|
||||
except ValueError:
|
||||
# V2 discovery emits one JSON event per line, rather than one envelope.
|
||||
result = {"events": [json.loads(line) for line in decoded.splitlines() if line.strip()]}
|
||||
if not isinstance(result, dict):
|
||||
raise ValueError
|
||||
except (UnicodeError, ValueError):
|
||||
return receipt(
|
||||
operation,
|
||||
{"returncode": process.returncode, "execution_status": "invalid_output_suppressed"},
|
||||
target=target,
|
||||
kind=kind,
|
||||
execution="unknown" if write else "completed",
|
||||
side_effect="possible" if write else "none",
|
||||
error_code="invalid_worker_output",
|
||||
files=files,
|
||||
)
|
||||
result = sanitized(result, secrets)
|
||||
if script == "evidence.py":
|
||||
result.pop("output", None)
|
||||
result["output_location"] = "artemis_task_attachments"
|
||||
status = result.get("execution_status", "")
|
||||
not_started = status in {
|
||||
"invalid_arguments",
|
||||
"tool_missing",
|
||||
"configuration_error",
|
||||
"launch_failed",
|
||||
}
|
||||
attempted = bool(
|
||||
result.get("write_executed")
|
||||
or result.get("write_may_have_started")
|
||||
or result.get("may_have_started")
|
||||
)
|
||||
verified_mismatch = (
|
||||
script == "device_adapter.py"
|
||||
and arguments[0] == "verify-version"
|
||||
and result.get("status") in {"FAIL", "INCOMPLETE"}
|
||||
)
|
||||
error = process.returncode != 0 and not verified_mismatch
|
||||
execution = (
|
||||
"not_started"
|
||||
if not_started
|
||||
else "unknown"
|
||||
if status in {"cleanup_unconfirmed", "unknown"}
|
||||
else "partial"
|
||||
if status in {"timed_out", "output_limit"} or (error and attempted)
|
||||
else "completed"
|
||||
)
|
||||
side_effect = (
|
||||
"confirmed"
|
||||
if write and result.get("write_executed")
|
||||
else "possible"
|
||||
if write and attempted
|
||||
else "none"
|
||||
)
|
||||
if write and script != "device_adapter.py" and not not_started:
|
||||
side_effect = "possible" if error else "confirmed"
|
||||
return receipt(
|
||||
operation,
|
||||
{"returncode": process.returncode, "tool_result": result, "stderr_suppressed": True},
|
||||
target=target,
|
||||
kind=kind,
|
||||
execution=execution,
|
||||
side_effect=side_effect,
|
||||
error_code=str(result.get("error_code") or "tool_execution_failed") if error else "",
|
||||
files=files,
|
||||
)
|
||||
Reference in new issue
Block a user