agnosis

tools/agnos/test_patch_system_reset_image.py

@@ -1,106 +1,88 @@
+import importlib.util
+import json
 from pathlib import Path
-import runpy
 
 import pytest
 
-from tools.agnos.patch_system_reset_image import (
-  AMDGPU_FIRMWARE_SHA256,
-  COMMA_SH_DISPLAY_WAIT_PATCH_MARKER,
-  comma_sh_has_expected_display_wait,
-  find_default_reference_manifest,
-  format_debugfs_mode,
-  patch_comma_sh_display_wait,
-  patch_setup_branding_script,
-  sha256_zstd_payload,
-)
+
+def _load_patch_module():
+  path = Path(__file__).resolve().parent / "patch_system_reset_image.py"
+  spec = importlib.util.spec_from_file_location("patch_system_reset_image_under_test", path)
+  module = importlib.util.module_from_spec(spec)
+  spec.loader.exec_module(module)
+  return module
 
 
-ORIGINAL_DISPLAY_WAIT = b'''#!/usr/bin/env bash
-echo "waiting for magic"
-for i in {1..200}; do
-  if systemctl is-active --quiet magic && [ -S /tmp/drmfd.sock ]; then
-    break
-  fi
-  sleep 0.1
-done
-
-if systemctl is-active --quiet magic && [ -S /tmp/drmfd.sock ]; then
-  echo "magic ready after ${SECONDS}s"
-else
-  echo "timed out waiting for magic, ${SECONDS}s"
-fi
-
-exec /data/continue.sh
-'''
+patch_image = _load_patch_module()
+ALLOWED_IMAGE_MUTATIONS = patch_image.ALLOWED_IMAGE_MUTATIONS
+UPSTREAM_PROTECTED_PAYLOADS = patch_image.UPSTREAM_PROTECTED_PAYLOADS
+VERSION_PATH_IN_IMAGE = patch_image.VERSION_PATH_IN_IMAGE
+resolve_upstream_manifest = patch_image.resolve_upstream_manifest
+update_manifest_system_entry = patch_image.update_manifest_system_entry
+write_allowed_file = patch_image.write_allowed_file
+expected_upstream_version = patch_image.expected_upstream_version
+verify_official_raw_hash = patch_image.verify_official_raw_hash
 
 
-def test_patch_comma_sh_display_wait_uses_available_display_service():
-  patched = patch_comma_sh_display_wait(ORIGINAL_DISPLAY_WAIT)
-
-  assert COMMA_SH_DISPLAY_WAIT_PATCH_MARKER.encode() in patched
-  assert b"systemctl cat magic.service" in patched
-  assert b"systemctl is-active --quiet magic" in patched
-  assert b"systemctl is-active --quiet weston-ready" in patched
-  assert b"[ -S /var/tmp/weston/wayland-0 ]" in patched
-  assert comma_sh_has_expected_display_wait(patched)
-  assert patch_comma_sh_display_wait(patched) == patched
+def test_only_version_is_mutable():
+  assert ALLOWED_IMAGE_MUTATIONS == {VERSION_PATH_IN_IMAGE}
+  assert set(UPSTREAM_PROTECTED_PAYLOADS) >= {
+    "/usr/comma/comma.sh",
+    "/usr/comma/reset",
+    "/usr/comma/setup",
+    "/usr/comma/updater",
+    "/etc/NetworkManager/NetworkManager.conf",
+    "/etc/NetworkManager/conf.d/10-globally-managed-devices.conf",
+    "/lib/systemd/system/NetworkManager.service",
+  }
 
 
-def test_patch_comma_sh_display_wait_rejects_unknown_layout():
-  with pytest.raises(RuntimeError, match="display readiness wait"):
-    patch_comma_sh_display_wait(b"#!/usr/bin/env bash\nexec /data/continue.sh\n")
+def test_starpilot_revision_maps_to_exact_upstream_version():
+  assert expected_upstream_version("19.6.3", None) == "19.6"
+  assert expected_upstream_version("19.6.3", "19.6-test") == "19.6-test"
+  with pytest.raises(RuntimeError, match="revision suffix"):
+    expected_upstream_version("19.6", None)
 
 
-@pytest.mark.parametrize("slider_text", ["slide to use", "slide to install"])
-def test_patch_mici_setup_branding_handles_old_and_new_labels(slider_text):
-  original = f'''OPENPILOT_URL = "https://openpilot.comma.ai"
-self._openpilot_slider = LargerSlider("{slider_text}\\nopenpilot", callback)
-self._continue_button = BigPillButton("install openpilot", green=True)
-self._continue_button.set_text("install openpilot" if not custom_software else "choose software")
-'''.encode()
+def test_official_source_hash_is_enforced(tmp_path):
+  image = tmp_path / "official.img"
+  image.write_bytes(b"official")
+  digest = patch_image.sha256_file(image)
 
-  patched = patch_setup_branding_script(original, "openpilot/system/ui/mici_setup.py")
-
-  assert b"installer.comma.ai/firestar5683/StarPilot" in patched
-  assert f"{slider_text}\\nstarpilot".encode() in patched
-  assert b"install StarPilot" in patched
-  assert b"install openpilot" not in patched
+  assert verify_official_raw_hash({"hash_raw": digest}, image) == digest
+  with pytest.raises(RuntimeError, match="hash mismatch"):
+    verify_official_raw_hash({"hash_raw": "0" * 64}, image)
 
 
-@pytest.mark.parametrize(("mode", "expected"), [
-  ("100775", "0100775"),
-  ("100644", "0100644"),
-  ("040755", "040755"),
-  ("120777", "0120777"),
-])
-def test_format_debugfs_mode(mode, expected):
-  assert format_debugfs_mode(mode) == expected
+def test_write_rejects_non_version_path(tmp_path):
+  with pytest.raises(RuntimeError, match="Refusing non-C3 AGNOS system mutation"):
+    write_allowed_file("debugfs", tmp_path / "system.img", "/usr/comma/setup", tmp_path / "setup")
 
 
-def test_external_gpu_firmware_matches_tinygrad_requirements():
-  firmware_metadata = Path(__file__).resolve().parents[2] / "tinygrad/runtime/autogen/am/fw.py"
-  hashes = runpy.run_path(firmware_metadata)["hashes"]
-  expected = {filename.removesuffix(".zst"): digest for filename, digest in AMDGPU_FIRMWARE_SHA256.items()}
-  assert all(hashes[filename] == digest for filename, digest in expected.items())
-
-
-def test_zstd_payload_hash(tmp_path):
-  import hashlib
-  import zstandard
-
-  payload = b"external GPU firmware payload"
-  compressed = tmp_path / "firmware.bin.zst"
-  compressed.write_bytes(zstandard.ZstdCompressor().compress(payload))
-
-  assert sha256_zstd_payload(compressed) == hashlib.sha256(payload).hexdigest()
-
-
-def test_default_reference_manifest_uses_sibling_openpilot(tmp_path):
-  primary = tmp_path / "starpilot/system/hardware/tici/agnos.json"
-  reference = tmp_path / "openpilot/openpilot/system/hardware/tici/agnos.json"
-  primary.parent.mkdir(parents=True)
-  reference.parent.mkdir(parents=True)
+def test_upstream_manifest_cannot_be_primary(tmp_path):
+  primary = tmp_path / "agnos.json"
   primary.write_text("[]")
-  reference.write_text("[]")
+  with pytest.raises(RuntimeError, match="Refusing to use StarPilot"):
+    resolve_upstream_manifest(str(primary), primary)
 
-  assert find_default_reference_manifest(primary) == reference.resolve()
+
+def test_update_manifest_changes_only_system_entry():
+  original = [
+    {"name": "boot", "url": "custom-boot", "hash": "boot-hash"},
+    {"name": "system", "url": "official", "hash": "old", "alt": {"url": "old-alt"}},
+  ]
+  updated = update_manifest_system_entry(original, "hosted", "new-hash", 123)
+
+  assert updated[0] == original[0]
+  assert updated[1] == {
+    "name": "system",
+    "url": "hosted",
+    "hash": "new-hash",
+    "hash_raw": "new-hash",
+    "size": 123,
+    "sparse": False,
+    "full_check": False,
+    "has_ab": True,
+    "ondevice_hash": "new-hash",
+  }
+  assert json.dumps(original)