485 lines
17 KiB
Python
485 lines
17 KiB
Python
"""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,
|
|
)
|