129 lines
4.3 KiB
Python
129 lines
4.3 KiB
Python
"""A fixed checksum tool using the installed ARTEMIS authoring SDK."""
|
|
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import sys
|
|
from typing import Annotated
|
|
import uuid
|
|
|
|
from mcp.server.fastmcp import FastMCP
|
|
from mcp.types import CallToolResult, TextContent, ToolAnnotations
|
|
from pydantic import Field
|
|
|
|
from artemis.plugins.sdk import ManagedProcessError, run_process, snapshot_file
|
|
|
|
mcp = FastMCP("managed-tool-demo")
|
|
CHECKSUM = "import hashlib,sys; print(hashlib.sha256(sys.stdin.buffer.read()).hexdigest())"
|
|
|
|
|
|
def failure(code: str, *, execution: str = "not_started") -> CallToolResult:
|
|
return CallToolResult(
|
|
isError=True,
|
|
content=[TextContent(type="text", text=code)],
|
|
structuredContent={
|
|
"artemis": {
|
|
"version": 1,
|
|
"execution": execution,
|
|
"side_effect": "none",
|
|
"error_code": code,
|
|
"retryable": False,
|
|
}
|
|
},
|
|
)
|
|
|
|
|
|
@mcp.tool(
|
|
annotations=ToolAnnotations(readOnlyHint=True, destructiveHint=False, idempotentHint=False)
|
|
)
|
|
async def verify_note(
|
|
name: Annotated[str, Field(min_length=1, max_length=180)],
|
|
sha256: Annotated[str, Field(pattern=r"^[a-fA-F0-9]{64}$")],
|
|
) -> CallToolResult:
|
|
"""Verify a task-owned file of at most 64 MiB against its explicitly approved SHA256."""
|
|
location = os.environ.get("ARTEMIS_ARTIFACTS_DIR")
|
|
if (
|
|
not location
|
|
or name in {".", ".."}
|
|
or Path(name).name != name
|
|
or any(c in name for c in '/\\:*?"<>|')
|
|
or name.rstrip(" .") != name
|
|
or any(ord(c) < 32 for c in name)
|
|
):
|
|
return failure("invalid_task_filename")
|
|
root = Path(location)
|
|
try:
|
|
if root.is_symlink() or root.is_junction() or root.resolve() != root.absolute():
|
|
return failure("invalid_task_directory")
|
|
snapshot = snapshot_file(root / name, sha256)
|
|
except (OSError, ValueError):
|
|
return failure("input_not_bound_to_approved_digest")
|
|
try:
|
|
# Never reopen the original pathname after verification. The fixed tool
|
|
# receives the same immutable bytes even if the source file changes.
|
|
result = await run_process(
|
|
[sys.executable, "-I", "-B", "-c", CHECKSUM],
|
|
cwd=root,
|
|
timeout=5,
|
|
input_data=snapshot.data,
|
|
)
|
|
except ManagedProcessError as exc:
|
|
return failure(
|
|
exc.error_code,
|
|
execution="unknown"
|
|
if exc.error_code == "process_cleanup_unconfirmed"
|
|
else "partial"
|
|
if exc.may_have_started
|
|
else "not_started",
|
|
)
|
|
if (
|
|
result.returncode != 0
|
|
or result.stdout.strip().decode("ascii", errors="replace") != snapshot.sha256
|
|
):
|
|
return failure("checksum_tool_result_invalid", execution="partial")
|
|
content = json.dumps(
|
|
{
|
|
"source_name": name,
|
|
"snapshot_sha256": snapshot.sha256,
|
|
"snapshot_bytes": snapshot.size,
|
|
"tool_returncode": result.returncode,
|
|
"test_verdict": "not_evaluated",
|
|
},
|
|
ensure_ascii=False,
|
|
).encode("utf-8")
|
|
receipt_name = f"verified-{uuid.uuid4().hex}.json"
|
|
try:
|
|
with (root / receipt_name).open("xb") as handle:
|
|
handle.write(content)
|
|
except OSError:
|
|
return failure("receipt_write_failed", execution="partial")
|
|
return CallToolResult(
|
|
content=[TextContent(type="text", text="已核对批准的文件快照，回执已保存。")],
|
|
structuredContent={
|
|
"snapshot_sha256": snapshot.sha256,
|
|
"artemis": {
|
|
"version": 1,
|
|
"execution": "completed",
|
|
"side_effect": "none",
|
|
"retryable": False,
|
|
"source": {
|
|
"kind": "file",
|
|
"target_id": "task-artifacts",
|
|
"label": "已批准文件快照",
|
|
},
|
|
"attachments": [
|
|
{
|
|
"name": receipt_name,
|
|
"media_type": "application/json",
|
|
"sha256": hashlib.sha256(content).hexdigest(),
|
|
}
|
|
],
|
|
},
|
|
},
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
mcp.run(transport="stdio")
|