awoooi/scripts/security/tests/test_runtime_image_mirror_controller.py

from __future__ import annotations

import importlib.util
import json
import sys
import tarfile
import tempfile
import unittest
from io import BytesIO
from pathlib import Path


ROOT = Path(__file__).resolve().parents[3]
CONTROLLER_PATH = ROOT / "scripts/security/runtime_image_mirror_controller.py"
POLICY_PATH = ROOT / "config/security/runtime-image-mirror-policy.json"


def _load_controller():
    spec = importlib.util.spec_from_file_location(
        "runtime_image_mirror_controller", CONTROLLER_PATH
    )
    assert spec is not None and spec.loader is not None
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


controller = _load_controller()


class RuntimeImageMirrorControllerTest(unittest.TestCase):
    def test_policy_is_internal_digest_pinned_and_runtime_cache_only(self) -> None:
        policy = controller.load_policy(POLICY_PATH)
        raw = POLICY_PATH.read_text(encoding="utf-8")

        self.assertEqual(policy.work_item_id, "AIA-P0-006-02A")
        self.assertEqual(policy.risk_level, "high")
        self.assertEqual(len(policy.images), 5)
        self.assertNotIn("github.com", raw)
        self.assertNotIn("ghcr.io", raw)
        self.assertIn('"external_pull_allowed": false', raw)
        self.assertTrue(
            all(
                image.runtime_ref.startswith(
                    "192.168.0.110:5000/awoooi/runtime-mirror/"
                )
                for image in policy.images
            )
        )
        self.assertTrue(
            all(
                image.runtime_ref.endswith("@" + image.target_registry_digest)
                for image in policy.images
            )
        )
        self.assertTrue(
            all(
                image.platform_digest != image.target_registry_digest
                for image in policy.images
            )
        )

    def test_policy_rejects_external_pull_and_mutable_runtime_ref(self) -> None:
        payload = json.loads(POLICY_PATH.read_text(encoding="utf-8"))
        with tempfile.TemporaryDirectory() as temp:
            path = Path(temp) / "policy.json"
            payload["source_contract"]["external_pull_allowed"] = True
            path.write_text(json.dumps(payload), encoding="utf-8")
            with self.assertRaisesRegex(
                controller.ControllerError, "external_pull_must_be_disabled"
            ):
                controller.load_policy(path)

            payload["source_contract"]["external_pull_allowed"] = False
            payload["images"][0][
                "runtime_ref"
            ] = "192.168.0.110:5000/awoooi/runtime-mirror/dex:latest"
            path.write_text(json.dumps(payload), encoding="utf-8")
            with self.assertRaisesRegex(
                controller.ControllerError, "runtime_ref_digest_mismatch"
            ):
                controller.load_policy(path)

    def test_archive_verifier_requires_source_platform_manifest(self) -> None:
        policy = controller.load_policy(POLICY_PATH)
        image = policy.images[0]
        with tempfile.TemporaryDirectory() as temp:
            archive_path = Path(temp) / "image.tar"
            member = "blobs/sha256/" + image.platform_digest.removeprefix("sha256:")
            with tarfile.open(archive_path, "w") as archive:
                info = tarfile.TarInfo(member)
                info.size = 2
                archive.addfile(info, BytesIO(b"{}"))

            self.assertTrue(
                controller._archive_contains_digest(archive_path, image.platform_digest)
            )
            self.assertFalse(
                controller._archive_contains_digest(
                    archive_path, image.target_registry_digest
                )
            )

    def test_source_manifests_match_internal_policy(self) -> None:
        policy = controller.load_policy(POLICY_PATH)
        by_asset = {image.asset_id: image for image in policy.images}
        kured = (ROOT / "k8s/kured/kured.yaml").read_text(encoding="utf-8")
        event = (ROOT / "k8s/observability/event-exporter.yaml").read_text(
            encoding="utf-8"
        )
        kube_state_metrics = (
            ROOT / "k8s/kube-state-metrics/kube-state-metrics.yaml"
        ).read_text(encoding="utf-8")
        node_problem_detector = (ROOT / "k8s/npd/node-problem-detector.yaml").read_text(
            encoding="utf-8"
        )

        self.assertIn(by_asset["runtime-image:kured"].runtime_ref, kured)
        self.assertIn(by_asset["runtime-image:event-exporter"].runtime_ref, event)
        self.assertIn(
            by_asset["runtime-image:kube-state-metrics"].runtime_ref,
            kube_state_metrics,
        )
        self.assertIn(
            by_asset["runtime-image:node-problem-detector"].runtime_ref,
            node_problem_detector,
        )
        self.assertNotIn("ghcr.io", kured)
        self.assertNotIn("ghcr.io", event)
        self.assertNotIn("registry.k8s.io", kube_state_metrics)
        self.assertNotIn("registry.k8s.io", node_problem_detector)

    def test_controller_and_cd_never_pull_external_images(self) -> None:
        controller_source = CONTROLLER_PATH.read_text(encoding="utf-8")
        workflow = (ROOT / ".gitea/workflows/cd.yaml").read_text(encoding="utf-8")

        self.assertNotIn("docker pull", controller_source)
        self.assertNotIn('"images", "pull"', controller_source)
        self.assertIn("Mirror Runtime Cache Images to Harbor", workflow)
        self.assertIn("Apply Runtime Image Mirror Policy", workflow)
        self.assertIn("runtime_image_mirror_controller.py", workflow)
        mirror_block = workflow[
            workflow.index(
                "- name: Mirror Runtime Cache Images to Harbor"
            ) : workflow.index("# ── API 鏡像建置")
        ]
        self.assertNotIn("HARBOR_PASSWORD", mirror_block)

    def test_verified_mirror_receipt_is_bound_to_trace_and_policy(self) -> None:
        policy = controller.load_policy(POLICY_PATH)
        receipt = {
            "schema_version": controller.RECEIPT_SCHEMA_VERSION,
            "trace_id": "aia-p0-006-02a-test",
            "work_item_id": policy.work_item_id,
            "policy_checksum": policy.checksum,
            "state": "verified",
            "verified_count": len(policy.images),
        }
        with tempfile.TemporaryDirectory() as temp:
            path = Path(temp) / "receipt.json"
            path.write_text(json.dumps(receipt), encoding="utf-8")
            self.assertEqual(
                controller._read_mirror_receipt(path, policy, "aia-p0-006-02a-test"),
                receipt,
            )
            with self.assertRaisesRegex(
                controller.ControllerError, "mirror_receipt_not_verified"
            ):
                controller._read_mirror_receipt(path, policy, "different-trace")


if __name__ == "__main__":
    unittest.main()