264 lines
11 KiB
Python
264 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import subprocess
|
|
import tempfile
|
|
import unittest
|
|
from unittest.mock import patch
|
|
|
|
from pi_dictate_service.config import ConfigError
|
|
from pi_dictate_service.startup import validate_startup
|
|
|
|
|
|
REPOSITORY_ROOT = Path(__file__).resolve().parents[3]
|
|
SCRIPT = REPOSITORY_ROOT / "services/dictate/install.sh"
|
|
|
|
|
|
class SystemdProfileTests(unittest.TestCase):
|
|
def run_install(self, home: Path, fake_bin: Path, *arguments: str) -> subprocess.CompletedProcess[str]:
|
|
environment = {
|
|
**os.environ,
|
|
"HOME": str(home),
|
|
"XDG_CONFIG_HOME": str(home / ".config"),
|
|
"PATH": f"{fake_bin}:{os.environ['PATH']}",
|
|
}
|
|
return subprocess.run(
|
|
[str(SCRIPT), *arguments],
|
|
env=environment,
|
|
capture_output=True,
|
|
text=True,
|
|
check=False,
|
|
)
|
|
|
|
def make_fake_bin(self, root: Path) -> Path:
|
|
fake_bin = root / "bin"
|
|
fake_bin.mkdir()
|
|
uv = fake_bin / "uv"
|
|
# Match uv sync's contract so transactional environment swapping is
|
|
# exercised without installing dependencies.
|
|
uv.write_text(
|
|
"#!/bin/sh\nmkdir -p \"$UV_PROJECT_ENVIRONMENT\"\nexit 0\n",
|
|
encoding="utf-8",
|
|
)
|
|
uv.chmod(0o755)
|
|
systemctl = fake_bin / "systemctl"
|
|
systemctl.write_text("#!/bin/sh\nexit 0\n", encoding="utf-8")
|
|
systemctl.chmod(0o755)
|
|
return fake_bin
|
|
|
|
def test_generated_cloud_dropin_replaces_offline_network_restrictions(self) -> None:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
home = root / "home"
|
|
home.mkdir()
|
|
fake_bin = self.make_fake_bin(root)
|
|
key = root / "key"
|
|
key.write_text("test-key", encoding="utf-8")
|
|
|
|
result = self.run_install(
|
|
home, fake_bin, "install", "--backend", "openai", "--api-key-file", str(key)
|
|
)
|
|
|
|
self.assertEqual(result.returncode, 0, result.stderr)
|
|
dropin = (
|
|
home / ".config/systemd/user/pi-dictate.service.d/"
|
|
"10-pi-dictate-network.conf"
|
|
).read_text(encoding="utf-8")
|
|
self.assertIn("Environment=PI_DICTATE_NETWORK_MODE=cloud", dropin)
|
|
self.assertIn(
|
|
"RestrictAddressFamilies=AF_UNIX AF_INET AF_INET6", dropin
|
|
)
|
|
self.assertIn("IPAddressDeny=", dropin)
|
|
self.assertNotIn("IPAddressDeny=any", dropin)
|
|
self.assertNotIn("HF_HUB_OFFLINE=1", dropin)
|
|
self.assertIn(
|
|
f'BindReadOnlyPaths="{home}/.config/pi-dictate/credentials/openai-api-key"',
|
|
dropin,
|
|
)
|
|
self.assertIn(
|
|
f'BindReadOnlyPaths="{home}/.local/share/pi-dictate/config/coding-terms.txt"',
|
|
dropin,
|
|
)
|
|
|
|
def test_check_rejects_tampered_network_dropin(self) -> None:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
home = root / "home"
|
|
home.mkdir()
|
|
fake_bin = self.make_fake_bin(root)
|
|
key = root / "key"
|
|
key.write_text("test-key", encoding="utf-8")
|
|
result = self.run_install(
|
|
home, fake_bin, "install", "--backend", "meta", "--api-key-file", str(key)
|
|
)
|
|
self.assertEqual(result.returncode, 0, result.stderr)
|
|
dropin = (
|
|
home / ".config/systemd/user/pi-dictate.service.d/"
|
|
"10-pi-dictate-network.conf"
|
|
)
|
|
with dropin.open("a", encoding="utf-8") as stream:
|
|
stream.write("IPAddressAllow=any\\n")
|
|
|
|
result = self.run_install(home, fake_bin, "check")
|
|
|
|
self.assertNotEqual(result.returncode, 0)
|
|
self.assertIn("expected managed policy", result.stderr)
|
|
|
|
def test_startup_requires_the_exact_validation_command(self) -> None:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
home = root / "home"
|
|
home.mkdir()
|
|
fake_bin = self.make_fake_bin(root)
|
|
key = root / "key"
|
|
key.write_text("test-key", encoding="utf-8")
|
|
result = self.run_install(
|
|
home, fake_bin, "install", "--backend", "meta", "--api-key-file", str(key)
|
|
)
|
|
self.assertEqual(result.returncode, 0, result.stderr)
|
|
environment = {
|
|
"HOME": str(home),
|
|
"XDG_CONFIG_HOME": str(home / ".config"),
|
|
"PI_DICTATE_CONFIG_PATH": str(home / ".config/pi-dictate/config.json"),
|
|
"PI_DICTATE_DEPENDENCY_PROFILE_PATH": str(
|
|
home / ".local/share/pi-dictate/.dependency-profile"
|
|
),
|
|
"PI_DICTATE_NETWORK_MODE": "cloud",
|
|
}
|
|
with patch.dict(os.environ, environment, clear=False):
|
|
validate_startup()
|
|
service = home / ".config/systemd/user/pi-dictate.service"
|
|
text = service.read_text(encoding="utf-8")
|
|
service.write_text(
|
|
"\n".join(
|
|
"ExecStartPre=/bin/true" if line.startswith("ExecStartPre=") else line
|
|
for line in text.splitlines()
|
|
) + "\n",
|
|
encoding="utf-8",
|
|
)
|
|
with self.assertRaisesRegex(ConfigError, "hardening"):
|
|
validate_startup()
|
|
|
|
def test_startup_rejects_a_world_writable_socket_unit(self) -> None:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
home = root / "home"
|
|
home.mkdir()
|
|
fake_bin = self.make_fake_bin(root)
|
|
key = root / "key"
|
|
key.write_text("test-key", encoding="utf-8")
|
|
result = self.run_install(
|
|
home, fake_bin, "install", "--backend", "meta", "--api-key-file", str(key)
|
|
)
|
|
self.assertEqual(result.returncode, 0, result.stderr)
|
|
environment = {
|
|
"HOME": str(home),
|
|
"XDG_CONFIG_HOME": str(home / ".config"),
|
|
"PI_DICTATE_CONFIG_PATH": str(home / ".config/pi-dictate/config.json"),
|
|
"PI_DICTATE_DEPENDENCY_PROFILE_PATH": str(
|
|
home / ".local/share/pi-dictate/.dependency-profile"
|
|
),
|
|
"PI_DICTATE_NETWORK_MODE": "cloud",
|
|
}
|
|
socket_unit = home / ".config/systemd/user/pi-dictate.socket"
|
|
socket_unit.write_text(
|
|
socket_unit.read_text(encoding="utf-8").replace(
|
|
"SocketMode=0600", "SocketMode=0666"
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
with patch.dict(os.environ, environment, clear=False):
|
|
with self.assertRaisesRegex(ConfigError, "socket systemd policy"):
|
|
validate_startup()
|
|
|
|
def test_local_dropin_binds_custom_model_and_dictionary_paths(self) -> None:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
home = root / "home"
|
|
home.mkdir()
|
|
fake_bin = self.make_fake_bin(root)
|
|
source = root / "source.nemo"
|
|
source.write_bytes(b"model")
|
|
digest = hashlib.sha256(source.read_bytes()).hexdigest()
|
|
result = self.run_install(
|
|
home,
|
|
fake_bin,
|
|
"install",
|
|
"--backend",
|
|
"local",
|
|
"--model-source",
|
|
str(source),
|
|
"--model-sha256",
|
|
digest,
|
|
)
|
|
self.assertEqual(result.returncode, 0, result.stderr)
|
|
custom_model = home / "custom model.nemo"
|
|
custom_terms = home / "custom terms.txt"
|
|
custom_model.write_bytes(source.read_bytes())
|
|
custom_terms.write_text("JSON\n", encoding="utf-8")
|
|
config_path = home / ".config/pi-dictate/config.json"
|
|
config = json.loads(config_path.read_text(encoding="utf-8"))
|
|
config["local"]["modelPath"] = str(custom_model)
|
|
config["local"]["termsPath"] = str(custom_terms)
|
|
config_path.write_text(json.dumps(config), encoding="utf-8")
|
|
config_path.chmod(0o600)
|
|
|
|
result = self.run_install(home, fake_bin, "configure", "--backend", "local")
|
|
|
|
self.assertEqual(result.returncode, 0, result.stderr)
|
|
dropin = (home / ".config/systemd/user/pi-dictate.service.d/10-pi-dictate-network.conf").read_text(encoding="utf-8")
|
|
self.assertIn(f'BindReadOnlyPaths="{custom_model}"', dropin)
|
|
self.assertIn(f'BindReadOnlyPaths="{custom_terms}"', dropin)
|
|
environment = {
|
|
"HOME": str(home),
|
|
"XDG_CONFIG_HOME": str(home / ".config"),
|
|
"PI_DICTATE_CONFIG_PATH": str(config_path),
|
|
"PI_DICTATE_DEPENDENCY_PROFILE_PATH": str(
|
|
home / ".local/share/pi-dictate/.dependency-profile"
|
|
),
|
|
"PI_DICTATE_NETWORK_MODE": "offline",
|
|
}
|
|
with patch.dict(os.environ, environment, clear=False):
|
|
validate_startup()
|
|
|
|
def test_local_dropin_and_profile_deny_are_restrictive(self) -> None:
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
root = Path(directory)
|
|
home = root / "home"
|
|
home.mkdir()
|
|
fake_bin = self.make_fake_bin(root)
|
|
model = root / "model"
|
|
model.write_bytes(b"model")
|
|
# Dry-run keeps this test independent of the pinned model bytes.
|
|
result = self.run_install(
|
|
home,
|
|
fake_bin,
|
|
"install",
|
|
"--backend",
|
|
"local",
|
|
"--model-source",
|
|
str(model),
|
|
"--model-sha256",
|
|
hashlib.sha256(model.read_bytes()).hexdigest(),
|
|
"--dry-run",
|
|
)
|
|
self.assertEqual(result.returncode, 0, result.stderr)
|
|
self.assertIn("offline network hardening", result.stdout)
|
|
|
|
profile = json.loads(
|
|
(REPOSITORY_ROOT / "profiles/nono/pi-local.json").read_text(
|
|
encoding="utf-8"
|
|
)
|
|
)
|
|
self.assertIn("$HOME/.local/share/pi-dictate", profile["filesystem"]["deny"])
|
|
# User decision: the profile is intentionally unchanged from HEAD.
|
|
self.assertNotIn("$HOME/.config/pi-dictate", profile["filesystem"]["deny"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|