300 lines
9.9 KiB
Python
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()
|