NumericalOS

tests/test_bootstrap_verify.py

back to source

import hashlib
import os
import subprocess
import tempfile
import unittest

from numos import seed
from numos.state import render

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
BOOTSTRAP = os.path.join(ROOT, "boot", "bootstrap.sh").replace("\\", "/")


def sh(snippet):
    full = 'NUMOS_SOURCE_ONLY=1 . "%s"; %s' % (BOOTSTRAP, snippet)
    env = dict(os.environ)
    env["NUMOS_SOURCE_ONLY"] = "1"
    env["NUMOS_LIB"] = os.path.dirname(BOOTSTRAP) + "/lib"
    proc = subprocess.run(["bash", "-c", full], capture_output=True, text=True, env=env)
    return proc.returncode, proc.stdout.strip(), proc.stderr.strip()


def write_temp(text):
    handle = tempfile.NamedTemporaryFile("w", suffix=".numos", delete=False, newline="\n")
    handle.write(text)
    handle.close()
    return handle.name.replace("\\", "/")


def good_state():
    return render({"version": 1, "arch_targets": seed.ARCH_TARGETS,
                   "boot_phases": seed.BOOT_PHASES, "units": seed.INFRA_UNITS,
                   "health": []})


class TestSha256(unittest.TestCase):
    def test_matches_python_hashlib(self):
        path = write_temp("hello\n")
        try:
            _, out, _ = sh('numos_sha256 "%s"' % path)
            self.assertEqual(out, hashlib.sha256(b"hello\n").hexdigest())
        finally:
            os.unlink(path)


class TestVerifySha256(unittest.TestCase):
    def test_matching_hash_succeeds_silently(self):
        path = write_temp("payload\n")
        want = hashlib.sha256(b"payload\n").hexdigest()
        try:
            code, _, err = sh('numos_verify_sha256 "%s" %s' % (path, want))
            self.assertEqual(code, 0, err)
        finally:
            os.unlink(path)

    def test_mismatched_hash_halts_with_both_values(self):
        path = write_temp("payload\n")
        try:
            code, _, err = sh('numos_verify_sha256 "%s" %s' % (path, "00" * 32))
            self.assertNotEqual(code, 0)
            self.assertIn("numos: HALT:", err)
            self.assertIn("hash mismatch", err)
        finally:
            os.unlink(path)

    def test_missing_file_halts(self):
        code, _, err = sh('numos_verify_sha256 "/nonexistent/path" %s' % ("00" * 32))
        self.assertNotEqual(code, 0)
        self.assertIn("numos: HALT:", err)


class TestVerifyState(unittest.TestCase):
    def test_good_state_verifies(self):
        path = write_temp(good_state())
        try:
            code, _, err = sh('numos_verify_state "%s"' % path)
            self.assertEqual(code, 0, err)
        finally:
            os.unlink(path)

    def test_tampered_state_halts(self):
        path = write_temp(good_state().replace("busybox-riscv64", "busybox-EVIL"))
        try:
            code, _, err = sh('numos_verify_state "%s"' % path)
            self.assertNotEqual(code, 0)
            self.assertIn("numos: HALT:", err)
            self.assertIn("state hash mismatch", err)
        finally:
            os.unlink(path)

    def test_state_without_c_record_halts(self):
        body = "\n".join(l for l in good_state().split("\n") if not l.startswith("C "))
        path = write_temp(body + "\n")
        try:
            code, _, err = sh('numos_verify_state "%s"' % path)
            self.assertNotEqual(code, 0)
            self.assertIn("numos: HALT:", err)
        finally:
            os.unlink(path)

    def test_state_with_multiple_c_records_halts(self):
        """State with two C records should be rejected with specific count."""
        body = good_state()
        # Insert a bogus second C record
        lines = body.split("\n")
        # Find the C record and add a duplicate
        c_line = [l for l in lines if l.startswith("C ")][0]
        lines.insert(1, c_line)  # Insert after the C record
        bad_state = "\n".join(lines)
        path = write_temp(bad_state)
        try:
            code, _, err = sh('numos_verify_state "%s"' % path)
            self.assertNotEqual(code, 0)
            self.assertIn("numos: HALT:", err)
            self.assertIn("2 C records", err)  # Verify the count is in the message
        finally:
            os.unlink(path)

    def test_temp_file_cleaned_up_on_halt_in_numos_sha256(self):
        """Temp file must be cleaned up even if numos_sha256 halts inside numos_verify_state.

        Uses a curated PATH with symlinks to necessary tools (grep, cut, head, rm) but
        excludes sha256sum and shasum, so numos_sha256 halts after temp file is created.
        Verifies trap cleanup removes the file before exit.
        """
        import shutil
        # Create a temporary directory for our curated PATH
        pathdir = tempfile.mkdtemp()
        tmpdir = tempfile.mkdtemp()
        try:
            # Create symlinks to essential tools in our curated PATH
            # These tools are needed by numos_verify_state to build the temp file body
            for tool in ["grep", "cut", "head", "rm"]:
                # Try to find the tool in standard locations
                tool_path = shutil.which(tool)
                if tool_path:
                    link_path = os.path.join(pathdir, tool)
                    # On Windows, symlinks might not work; use wrapper script instead
                    try:
                        os.symlink(tool_path, link_path)
                    except (OSError, NotImplementedError):
                        # Fallback: create wrapper script with properly quoted path
                        with open(link_path, "w") as f:
                            # Properly quote the path in the wrapper script
                            f.write(f'#!/bin/bash\nexec "{tool_path}" "$@"\n')
                        os.chmod(link_path, 0o755)

            # Create a good state file
            state_path = write_temp(good_state())
            try:
                # Set up the shell invocation with curated PATH
                # Hash tools are not in PATH, so numos_sha256 will fail
                full = 'NUMOS_SOURCE_ONLY=1 . "%s"; numos_verify_state "%s"' % (BOOTSTRAP, state_path)
                env = dict(os.environ)
                env["NUMOS_SOURCE_ONLY"] = "1"
                env["NUMOS_LIB"] = os.path.dirname(BOOTSTRAP) + "/lib"
                env["TMPDIR"] = tmpdir
                env["PATH"] = pathdir  # Only our curated bin directory

                proc = subprocess.run(["bash", "-c", full], capture_output=True, text=True, env=env)

                # Verify the command failed (numos_sha256 halts due to missing hash tools)
                self.assertNotEqual(proc.returncode, 0)
                self.assertIn("numos: HALT:", proc.stderr)

                # Verify temp file was cleaned up by trap (should be empty or no numos-verify files)
                remaining_files = [f for f in os.listdir(tmpdir) if f.startswith("numos-verify")]
                self.assertEqual(len(remaining_files), 0,
                               f"Temp files not cleaned up: {remaining_files}")
            finally:
                os.unlink(state_path)
        finally:
            # Clean up curated PATH directory
            for f in os.listdir(pathdir):
                try:
                    os.unlink(os.path.join(pathdir, f))
                except OSError:
                    pass
            os.rmdir(pathdir)
            os.rmdir(tmpdir)


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