cang-ying/remote-node/tests/test_end_to_end.py

300 lines
9.9 KiB
Python

from __future__ import annotations
import importlib.util
import json
import os
import sys
import tempfile
import threading
import unittest
from pathlib import Path
from urllib.error import HTTPError
from urllib.parse import urlencode
from urllib.request import Request, urlopen
ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
spec = importlib.util.spec_from_file_location(
"control_plane", ROOT / "server" / "control_plane.py"
)
control_plane = importlib.util.module_from_spec(spec)
assert spec.loader
spec.loader.exec_module(control_plane)
class FlowTest(unittest.TestCase):
def setUp(self) -> None:
self.temp = tempfile.TemporaryDirectory()
self.admin = "admin-" + "a" * 32
self.owner = "owner-" + "o" * 32
self.persona = "persona-" + "p" * 32
self.operator = "operator-" + "z" * 32
app = control_plane.ControlPlane(
os.path.join(self.temp.name, "control.sqlite3"),
self.admin,
self.owner,
self.persona,
self.operator,
"signing-" + "s" * 32,
)
self.app = app
self.server = control_plane.AppServer(("127.0.0.1", 0), app)
self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
self.thread.start()
self.base = f"http://127.0.0.1:{self.server.server_port}"
def tearDown(self) -> None:
self.server.shutdown()
self.server.server_close()
self.temp.cleanup()
def api(self, method, path, token=None, body=None):
data = None if body is None else json.dumps(body).encode()
headers = {"Accept": "application/json"}
if token:
headers["Authorization"] = f"Bearer {token}"
if data is not None:
headers["Content-Type"] = "application/json"
with urlopen(
Request(self.base + path, data=data, headers=headers, method=method)
) as response:
return json.load(response)
def test_pair_approve_execute_receipt(self) -> None:
pairing = self.api("POST", "/v1/pairings", self.admin, {})
node = self.api(
"POST",
"/v1/pairings/claim",
body={
"code": pairing["code"],
"node_name": "canger-home-ubuntu",
"agent_version": "test",
"policy": {
"allowed_roots": ["/srv/canger"],
"enabled_actions": [
"system.status",
"repo.status",
"command.run",
],
},
},
)
with self.assertRaises(HTTPError) as reused:
self.api(
"POST",
"/v1/pairings/claim",
body={
"code": pairing["code"],
"node_name": "attacker",
"policy": {
"allowed_roots": ["/srv/canger"],
"enabled_actions": ["system.status"],
},
},
)
self.assertEqual(reused.exception.code, 401)
created = self.api(
"POST",
"/v1/tasks",
self.operator,
{
"node_id": node["node_id"],
"action": "system.status",
"args": {},
"requester_id": "ICE-GL-ZY001",
"requester_name": "铸渊",
},
)
waiting = self.api(
"GET",
f"/v1/nodes/{node['node_id']}/tasks/next",
node["node_token"],
)
self.assertIsNone(waiting["task"])
self.api(
"POST",
f"/v1/tasks/{created['task_id']}/approve",
self.owner,
{},
)
leased = self.api(
"GET",
f"/v1/nodes/{node['node_id']}/tasks/next",
node["node_token"],
)
self.assertEqual(leased["task"]["action"], "system.status")
event = self.api(
"POST",
f"/v1/tasks/{created['task_id']}/events",
node["node_token"],
{"kind": "started", "payload": {"agent_version": "test"}},
)
self.assertGreater(event["seq"], 0)
events = self.api(
"GET",
f"/v1/tasks/{created['task_id']}/events?after=0",
self.persona,
)
self.assertEqual(events["events"][0]["kind"], "started")
self.api(
"POST",
f"/v1/tasks/{created['task_id']}/result",
node["node_token"],
{"ok": True, "output": "healthy"},
)
receipt = self.api(
"GET", f"/v1/tasks/{created['task_id']}", self.operator
)
self.assertEqual(receipt["status"], "completed")
self.assertEqual(receipt["result"]["output"], "healthy")
def test_time_limited_grant_auto_approves_safe_task(self) -> None:
pairing = self.api("POST", "/v1/pairings", self.admin, {})
node = self.api(
"POST",
"/v1/pairings/claim",
body={
"code": pairing["code"],
"node_name": "canger-grant-test",
"agent_version": "test",
"policy": {
"allowed_roots": ["/srv/canger"],
"enabled_actions": ["system.status", "repo.status"],
},
},
)
grant = self.api(
"POST",
"/v1/grants",
self.operator,
{
"node_id": node["node_id"],
"duration_seconds": 3600,
"reason": "测试标准开发会话",
"requester_id": "ICE-GL-ZY001",
"requester_name": "铸渊",
},
)
link_expiry = control_plane.now_ts() + 1800
token = self.app.grant_token(grant["grant_id"], link_expiry)
form = urlencode({"id": grant["grant_id"], "token": token}).encode()
request = Request(
self.base + "/approve/grant",
data=form,
headers={"Content-Type": "application/x-www-form-urlencoded"},
method="POST",
)
with urlopen(request) as response:
self.assertEqual(response.status, 200)
task = self.api(
"POST",
"/v1/tasks",
self.operator,
{
"node_id": node["node_id"],
"action": "repo.status",
"args": {"repo": "/srv/canger/project"},
"requester_id": "ICE-GL-ZY001",
"requester_name": "铸渊",
},
)
self.assertEqual(task["status"], "approved")
self.assertFalse(task["approval_required"])
self.assertEqual(task["grant_id"], grant["grant_id"])
def test_persona_can_request_but_cannot_approve_its_own_grant(self) -> None:
pairing = self.api("POST", "/v1/pairings", self.admin, {})
node = self.api(
"POST",
"/v1/pairings/claim",
body={
"code": pairing["code"],
"node_name": "canger-persona-request-test",
"agent_version": "test",
"policy": {
"allowed_roots": ["/srv/canger"],
"enabled_actions": ["system.status", "repo.status"],
},
},
)
grant = self.api(
"POST",
"/v1/grants",
self.persona,
{
"node_id": node["node_id"],
"duration_seconds": 4 * 3600,
"reason": "检查苍耳仓库状态",
"requester_id": "PTS-VA-001-EED",
"requester_name": "耳耳蛋",
},
)
self.assertEqual(grant["status"], "pending")
self.assertEqual(grant["requested_by"], "PTS-VA-001-EED")
self.assertEqual(grant["requester_name"], "耳耳蛋")
self.assertEqual(grant["notification"], "not_configured")
self.assertEqual(grant["roots"], ["/srv/canger"])
with self.assertRaises(HTTPError) as unauthorized:
self.api(
"POST",
f"/v1/tasks/{grant['grant_id']}/approve",
self.persona,
{},
)
self.assertEqual(unauthorized.exception.code, 401)
pending_task = self.api(
"POST",
"/v1/tasks",
self.persona,
{
"node_id": node["node_id"],
"action": "system.status",
"args": {},
"requester_id": "PTS-VA-001-EED",
"requester_name": "耳耳蛋",
},
)
self.assertEqual(pending_task["status"], "pending")
self.assertEqual(pending_task["requested_by"], "PTS-VA-001-EED")
link_expiry = control_plane.now_ts() + 1800
token = self.app.grant_token(grant["grant_id"], link_expiry)
form = urlencode({"id": grant["grant_id"], "token": token}).encode()
approval = Request(
self.base + "/approve/grant",
data=form,
headers={"Content-Type": "application/x-www-form-urlencoded"},
method="POST",
)
with urlopen(approval) as response:
self.assertEqual(response.status, 200)
approved_task = self.api(
"POST",
"/v1/tasks",
self.persona,
{
"node_id": node["node_id"],
"action": "repo.status",
"args": {"repo": "/srv/canger/project"},
"requester_id": "PTS-VA-001-EED",
"requester_name": "耳耳蛋",
},
)
self.assertEqual(approved_task["status"], "approved")
self.assertEqual(approved_task["requested_by"], "PTS-VA-001-EED")
self.assertEqual(approved_task["grant_id"], grant["grant_id"])
if __name__ == "__main__":
unittest.main()