215 lines
10 KiB
Python
215 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import copy
|
|
import hashlib
|
|
import json
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from .policy import (
|
|
CompositionError,
|
|
PROVIDER_CREDENTIAL_DEFINITIONS,
|
|
PROVIDER_PROFILE_NAMES,
|
|
RAW_CREDENTIAL_VARIABLES,
|
|
assert_single_lm_studio_registration,
|
|
compose_profile,
|
|
isolated_settings,
|
|
load_profile,
|
|
)
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[4]
|
|
PROFILES = ROOT / "profiles" / "nono"
|
|
LM_STUDIO = ROOT / "extensions" / "lm-studio" / "index.ts"
|
|
PI_INSTALL = "/home/jb/.bun/install/global/node_modules/@earendil-works/pi-coding-agent"
|
|
JITI_INSTALL = "/home/jb/.bun/install/global/node_modules/jiti"
|
|
GLOBAL_INSTRUCTIONS = "/home/jb/.pi/agent/AGENTS.md"
|
|
|
|
|
|
class PolicyContractTests(unittest.TestCase):
|
|
def test_cores_are_fail_closed_and_separate_state(self) -> None:
|
|
for context in ("personal", "work"):
|
|
with self.subTest(context=context):
|
|
profile = load_profile(PROFILES / f"pi-core-{context}.json")
|
|
self.assertEqual(profile["extends"], "default")
|
|
self.assertEqual(profile["workdir"], {"access": "none"})
|
|
self.assertEqual(profile["network"], {"block": True})
|
|
self.assertEqual(
|
|
set(profile["environment"]["deny_vars"]),
|
|
RAW_CREDENTIAL_VARIABLES,
|
|
)
|
|
self.assertNotIn("SSH_AUTH_SOCK", profile["environment"]["allow_vars"])
|
|
self.assertIn("system_write_linux", profile["groups"]["exclude"])
|
|
self.assertIn("user_tools", profile["groups"]["exclude"])
|
|
grants = json.dumps(profile["filesystem"], sort_keys=True)
|
|
self.assertIn(f"pi-sbx/{context}", grants)
|
|
other = "work" if context == "personal" else "personal"
|
|
self.assertNotIn(f"pi-sbx/{other}", grants)
|
|
self.assertNotIn("$WORKDIR", grants)
|
|
self.assertNotIn('"/tmp"', grants)
|
|
self.assertNotIn("$TMPDIR", grants)
|
|
self.assertEqual(
|
|
set(profile["filesystem"]["allow"]),
|
|
{
|
|
f"$XDG_CACHE_HOME/pi-sbx/{context}",
|
|
f"$XDG_STATE_HOME/pi-sbx/{context}/agent",
|
|
f"$XDG_STATE_HOME/pi-sbx/{context}/sessions",
|
|
},
|
|
)
|
|
self.assertEqual(
|
|
profile["filesystem"]["read"],
|
|
[PI_INSTALL, JITI_INSTALL],
|
|
)
|
|
self.assertEqual(
|
|
set(profile["filesystem"]["read_file"]),
|
|
{GLOBAL_INSTRUCTIONS, "/usr/bin/node"},
|
|
)
|
|
self.assertEqual(
|
|
set(profile["filesystem"]["allow_file"]),
|
|
{"/dev/null", "/dev/stderr", "/dev/stdout", "/dev/tty"},
|
|
)
|
|
for capability in ("write", "write_file", "unix_socket"):
|
|
self.assertEqual(profile["filesystem"].get(capability, []), [])
|
|
|
|
def test_provider_profiles_select_one_route(self) -> None:
|
|
for provider, profile_name in PROVIDER_PROFILE_NAMES.items():
|
|
with self.subTest(provider=provider):
|
|
profile = load_profile(PROFILES / f"{profile_name}.json")
|
|
network = profile["network"]
|
|
self.assertFalse(network["block"])
|
|
if provider == "lm-studio":
|
|
self.assertEqual(set(network), {"block", "connect_port"})
|
|
self.assertEqual(network["connect_port"], [1234])
|
|
self.assertNotIn("credentials", network)
|
|
else:
|
|
self.assertEqual(
|
|
set(network),
|
|
{"block", "allow_domain", "credentials", "custom_credentials"},
|
|
)
|
|
self.assertEqual(network["allow_domain"], [])
|
|
self.assertEqual(network["credentials"], [provider])
|
|
self.assertEqual(
|
|
network["custom_credentials"],
|
|
{provider: PROVIDER_CREDENTIAL_DEFINITIONS[provider]},
|
|
)
|
|
self.assertNotIn("connect_port", network)
|
|
self.assertEqual(profile["workdir"], {"access": "none"})
|
|
self.assertIn("system_read_linux_core", profile["groups"]["exclude"])
|
|
self.assertIn("system_read_macos", profile["groups"]["exclude"])
|
|
self.assertTrue(profile["environment"]["allow_vars"])
|
|
self.assertEqual(
|
|
set(profile["environment"]["deny_vars"]),
|
|
RAW_CREDENTIAL_VARIABLES,
|
|
)
|
|
self.assertEqual(profile["open_urls"]["allow_origins"], [])
|
|
self.assertFalse(profile["allow_launch_services"])
|
|
|
|
def test_trusted_composition_replaces_block_with_one_provider(self) -> None:
|
|
core = load_profile(PROFILES / "pi-core-personal.json")
|
|
provider = load_profile(PROFILES / "provider-lm-studio.json")
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
temporary_directory = Path(directory)
|
|
generated = compose_profile(
|
|
core,
|
|
provider,
|
|
workdir_access="read",
|
|
temporary_directory=temporary_directory,
|
|
lm_studio_extension=LM_STUDIO,
|
|
)
|
|
self.assertEqual(generated["workdir"], {"access": "read"})
|
|
self.assertIn(str(temporary_directory), generated["filesystem"]["allow"])
|
|
self.assertEqual(
|
|
generated["environment"]["set_vars"]["TMPDIR"],
|
|
str(temporary_directory),
|
|
)
|
|
self.assertEqual(generated["network"]["connect_port"], [1234])
|
|
self.assertFalse(generated["network"]["block"])
|
|
self.assertNotIn("command_policies", generated)
|
|
self.assertIn(str(LM_STUDIO), generated["filesystem"]["read"])
|
|
with self.assertRaisesRegex(CompositionError, "cannot load"):
|
|
compose_profile(
|
|
core,
|
|
load_profile(PROFILES / "provider-openai.json"),
|
|
workdir_access="read",
|
|
temporary_directory=temporary_directory,
|
|
lm_studio_extension=LM_STUDIO,
|
|
)
|
|
with self.assertRaisesRegex(CompositionError, "private owned"):
|
|
compose_profile(
|
|
core,
|
|
provider,
|
|
workdir_access="read",
|
|
temporary_directory=Path("/tmp"),
|
|
lm_studio_extension=LM_STUDIO,
|
|
)
|
|
mutated_core = copy.deepcopy(core)
|
|
mutated_core["filesystem"]["allow"].append("/tmp")
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
with self.assertRaisesRegex(CompositionError, "differs from reviewed"):
|
|
compose_profile(
|
|
mutated_core,
|
|
provider,
|
|
workdir_access="read",
|
|
temporary_directory=Path(directory),
|
|
lm_studio_extension=LM_STUDIO,
|
|
)
|
|
|
|
def test_generated_settings_have_no_privileged_resources(self) -> None:
|
|
normal = isolated_settings("openai")
|
|
self.assertEqual(normal["packages"], [])
|
|
self.assertEqual(normal["extensions"], [])
|
|
self.assertEqual(normal["defaultTools"], ["read", "edit", "write"])
|
|
self.assertNotIn("bash", normal["defaultTools"])
|
|
self.assertEqual(normal["defaultProjectTrust"], "never")
|
|
local = isolated_settings("lm-studio", lm_studio_extension=LM_STUDIO)
|
|
self.assertEqual(local["extensions"], [str(LM_STUDIO)])
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
unreviewed = Path(directory) / "lm-studio.ts"
|
|
unreviewed.write_text("unreviewed\n", encoding="utf-8")
|
|
with self.assertRaisesRegex(CompositionError, "reviewed lock source"):
|
|
isolated_settings("lm-studio", lm_studio_extension=unreviewed)
|
|
serialized = json.dumps(local).lower()
|
|
for forbidden in ("pi-workflow", "update", "background", "notify"):
|
|
self.assertNotIn(forbidden, serialized)
|
|
|
|
def test_lm_studio_registration_is_unique_and_uses_exact_ip(self) -> None:
|
|
lock = json.loads((ROOT / "config" / "pi-package-lock.json").read_text(encoding="utf-8"))
|
|
entry = next(item for item in lock["extensions"] if item["name"] == "lm-studio")
|
|
self.assertEqual(entry["source"], "extensions/lm-studio/index.ts")
|
|
self.assertEqual(entry["status"], "blocked")
|
|
self.assertEqual(
|
|
entry["reviewed"]["sha256"],
|
|
hashlib.sha256(LM_STUDIO.read_bytes()).hexdigest(),
|
|
)
|
|
assert_single_lm_studio_registration([LM_STUDIO])
|
|
contents = LM_STUDIO.read_text(encoding="utf-8")
|
|
self.assertIn('"http://127.0.0.1:1234/v1"', contents)
|
|
self.assertNotIn("http://localhost:1234", contents)
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
duplicate = Path(directory) / "duplicate.ts"
|
|
duplicate.write_text(
|
|
'pi.registerProvider("lm-studio", {});\n',
|
|
encoding="utf-8",
|
|
)
|
|
with self.assertRaisesRegex(CompositionError, "exactly one"):
|
|
assert_single_lm_studio_registration([LM_STUDIO, duplicate])
|
|
duplicate.write_text(
|
|
'pi.registerProvider("lm-studio", {});\n'
|
|
'pi.registerProvider("lm-studio", {});\n',
|
|
encoding="utf-8",
|
|
)
|
|
with self.assertRaisesRegex(CompositionError, "found 2"):
|
|
assert_single_lm_studio_registration([duplicate])
|
|
duplicate.write_text(
|
|
'const baseUrl = "http://localhost:1234/v1";\n'
|
|
'pi.registerProvider("lm-studio", {});\n',
|
|
encoding="utf-8",
|
|
)
|
|
with self.assertRaisesRegex(CompositionError, "legacy localhost"):
|
|
assert_single_lm_studio_registration([duplicate])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|