184 lines
8.0 KiB
Python
184 lines
8.0 KiB
Python
from __future__ import annotations
|
|
|
|
import io
|
|
import json
|
|
import os
|
|
import stat
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from agent_comm.server import (
|
|
MAX_BODY_BYTES,
|
|
MAX_FILE_BYTES,
|
|
MAX_RPC_BYTES,
|
|
SCHEMA_VERSION,
|
|
AgentCommServer,
|
|
MailboxError,
|
|
MessageError,
|
|
MessageStore,
|
|
SecretContentError,
|
|
redact_text,
|
|
utc_now,
|
|
)
|
|
|
|
|
|
class MessageStoreTests(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
self.temp_dir = tempfile.TemporaryDirectory()
|
|
self.root = Path(self.temp_dir.name) / "mailbox"
|
|
self.store = MessageStore(self.root)
|
|
|
|
def tearDown(self) -> None:
|
|
self.temp_dir.cleanup()
|
|
|
|
def message(self, *, sender: str = "oldbox", body: str = "sanitized report") -> dict[str, object]:
|
|
return {
|
|
"schema_version": SCHEMA_VERSION,
|
|
"id": "12345678-1234-4234-8234-123456789abc",
|
|
"created_at": utc_now(),
|
|
"sender": sender,
|
|
"kind": "report" if sender == "oldbox" else "instruction",
|
|
"correlation_id": None,
|
|
"title": "Test message",
|
|
"body": body,
|
|
"attachments": [],
|
|
}
|
|
|
|
def test_layout_is_private(self) -> None:
|
|
for path in [self.root, *(self.root / name for name in ("inbox", "outbox", "processed", "archive"))]:
|
|
self.assertTrue(path.is_dir())
|
|
self.assertEqual(stat.S_IMODE(path.stat().st_mode) & 0o077, 0)
|
|
|
|
def test_send_report_is_atomic_and_private(self) -> None:
|
|
result = self.store.send_report(title="Status", body="all clear")
|
|
target = Path(result["path"])
|
|
self.assertTrue(target.is_file())
|
|
self.assertEqual(stat.S_IMODE(target.stat().st_mode) & 0o077, 0)
|
|
self.assertEqual(json.loads(target.read_text(encoding="utf-8"))["body"], "all clear")
|
|
self.assertFalse(list(self.store.outbox.glob("*.tmp")))
|
|
|
|
def test_secret_pattern_is_refused_fail_closed(self) -> None:
|
|
with self.assertRaises(SecretContentError):
|
|
self.store.send_report(
|
|
title="Diagnostic",
|
|
body="password=example-only\nAuthorization: Bearer example-only\n{\"token\": \"example-only\"}\nnormal authentication prose",
|
|
)
|
|
redacted, changed = redact_text("Authorization: Basic example-only\nnormal authentication prose")
|
|
self.assertTrue(changed)
|
|
self.assertNotIn("example-only", redacted)
|
|
self.assertIn("normal authentication prose", redacted)
|
|
|
|
def test_redact_text_preserves_non_secret_prose(self) -> None:
|
|
value, changed = redact_text("The password authentication path is probably used.")
|
|
self.assertEqual(value, "The password authentication path is probably used.")
|
|
self.assertFalse(changed)
|
|
|
|
def test_size_limits_and_schema_validation(self) -> None:
|
|
with self.assertRaises(MessageError):
|
|
self.store.send_report(title="x", body="x" * (MAX_BODY_BYTES + 1))
|
|
invalid = self.message()
|
|
invalid["attachments"] = ["/etc/passwd"]
|
|
with self.assertRaises(MessageError):
|
|
self.store._write_message("outbox", invalid)
|
|
invalid = self.message()
|
|
invalid["id"] = "../escape"
|
|
with self.assertRaises(MessageError):
|
|
self.store._write_message("outbox", invalid)
|
|
invalid = self.message()
|
|
invalid["schema_version"] = 2
|
|
with self.assertRaises(MessageError):
|
|
self.store._write_message("outbox", invalid)
|
|
|
|
def test_inbox_messages_are_validated_and_moved(self) -> None:
|
|
message = self.message(sender="mac-agent", body="A sanitized instruction")
|
|
path = self.store.inbox / f"{message['id']}.json"
|
|
path.write_text(json.dumps(message), encoding="utf-8")
|
|
os.chmod(path, 0o600)
|
|
messages = self.store.read_instructions()
|
|
self.assertEqual(messages[0]["body"], "A sanitized instruction")
|
|
self.assertFalse(path.exists())
|
|
self.assertTrue((self.store.processed / path.name).is_file())
|
|
|
|
def test_incoming_secret_is_refused_and_removed(self) -> None:
|
|
message = self.message(sender="mac-agent", body="token=example-only")
|
|
path = self.store.inbox / f"{message['id']}.json"
|
|
path.write_text(json.dumps(message), encoding="utf-8")
|
|
os.chmod(path, 0o600)
|
|
self.assertEqual(self.store.read_instructions(), [])
|
|
self.assertFalse(path.exists())
|
|
self.assertFalse((self.store.processed / path.name).exists())
|
|
|
|
def test_symlink_paths_are_refused(self) -> None:
|
|
outside = Path(self.temp_dir.name) / "outside.json"
|
|
outside.write_text("{}", encoding="utf-8")
|
|
link = self.store.inbox / "12345678-1234-4234-8234-123456789abc.json"
|
|
os.symlink(outside, link)
|
|
self.assertEqual(self.store.read_instructions(), [])
|
|
self.assertTrue(outside.exists())
|
|
root_link = Path(self.temp_dir.name) / "root-link"
|
|
os.symlink(self.root, root_link)
|
|
with self.assertRaises(MailboxError):
|
|
MessageStore(root_link)
|
|
|
|
def test_check_mode_does_not_create_missing_root(self) -> None:
|
|
missing = Path(self.temp_dir.name) / "missing"
|
|
with self.assertRaises(MailboxError):
|
|
MessageStore(missing, create=False)
|
|
self.assertFalse(missing.exists())
|
|
|
|
def test_encoded_file_limit(self) -> None:
|
|
self.assertLess(MAX_FILE_BYTES, MAX_BODY_BYTES * 4)
|
|
path = self.store.inbox / "bad.json"
|
|
path.write_bytes(b"x" * (MAX_FILE_BYTES + 1))
|
|
os.chmod(path, 0o600)
|
|
self.assertEqual(self.store.read_instructions(), [])
|
|
|
|
def test_nas_mode_requires_confirmed_mount(self) -> None:
|
|
with self.assertRaises(MailboxError):
|
|
MessageStore(self.root, deployment_mode="nas", nas_mount="/definitely/not/a/mount", nas_confirmed=True)
|
|
|
|
def test_mcp_initialize_status_and_invalid_id(self) -> None:
|
|
server = AgentCommServer(self.store)
|
|
initialized = server.handle_json_rpc(
|
|
{"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {"protocolVersion": "2024-11-05"}}
|
|
)
|
|
self.assertEqual(initialized["result"]["serverInfo"]["name"], "agent-comm")
|
|
status = server.call_tool("agent_status", {})
|
|
self.assertFalse(json.loads(status["content"][0]["text"])["shared"])
|
|
invalid = server.handle_json_rpc({"jsonrpc": "2.0", "id": {"secret": "example-only"}, "method": "ping"})
|
|
self.assertIsNone(invalid["id"])
|
|
self.assertNotIn("example-only", json.dumps(invalid))
|
|
|
|
def test_stdio_input_is_bounded_and_malformed_requests_are_rejected(self) -> None:
|
|
server = AgentCommServer(self.store)
|
|
output = io.StringIO()
|
|
server.serve(io.StringIO("x" * (MAX_RPC_BYTES + 1) + "\n"), output)
|
|
response = json.loads(output.getvalue())
|
|
self.assertEqual(response["error"]["code"], -32600)
|
|
output = io.StringIO()
|
|
server.serve(io.StringIO('{"jsonrpc":"1.0","id":1,"method":"ping"}\n'), output)
|
|
self.assertEqual(json.loads(output.getvalue())["error"]["code"], -32600)
|
|
|
|
def test_send_ack_is_validated(self) -> None:
|
|
correlation_id = "12345678-1234-4234-8234-123456789abc"
|
|
result = self.store.send_ack(correlation_id=correlation_id, text="Received; no action taken.")
|
|
value = json.loads(Path(result["path"]).read_text(encoding="utf-8"))
|
|
self.assertEqual(value["kind"], "ack")
|
|
self.assertEqual(value["correlation_id"], correlation_id)
|
|
with self.assertRaises(SecretContentError):
|
|
self.store.send_ack(correlation_id=correlation_id, text="access_token=example-only")
|
|
|
|
def test_publication_refuses_to_overwrite_existing_message(self) -> None:
|
|
message = self.message()
|
|
self.store._write_message("outbox", message)
|
|
with self.assertRaises(MessageError):
|
|
self.store._write_message("outbox", message)
|
|
target = self.store.outbox / f"{message['id']}.json"
|
|
self.assertEqual(json.loads(target.read_text(encoding="utf-8"))["id"], message["id"])
|
|
self.assertFalse(list(self.store.outbox.glob("*.tmp")))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|