tai/taicpp/tests/seal_diff.py

#!/usr/bin/env python3
"""Differential seal test: tai.py (python) against build/seal_probe (C++), both directions.
Run from taicpp/: python3 tests/seal_diff.py"""
import json
import os
import sqlite3
import subprocess
import sys
import tempfile
import unittest
import warnings

warnings.simplefilter("ignore", ResourceWarning)

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
sys.path.insert(0, os.path.dirname(ROOT))
import tai  # noqa: E402

PROBE = os.path.join(ROOT, "build", "seal_probe")


def probe(*args):
    out = subprocess.run([PROBE] + list(args), capture_output=True, text=True, timeout=120).stdout
    return out.strip().splitlines()[-1] if out.strip() else ""


class Args:
    profile = "default"
    yes = True


class SealDiff(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.home = self.tmp.name
        os.environ["TAI_HOME"] = self.home
        os.environ.pop("TAI_PASSPHRASE", None)

    def tearDown(self):
        os.environ.pop("TAI_HOME", None)
        self.tmp.cleanup()

    def test_master_and_keystream(self):
        seal = tai.Seal(self.home, "pass-ä-1")
        self.assertEqual(probe("master", self.home, "pass-ä-1"), seal.master.hex())
        nonce = bytes(range(16))
        for length in (0, 1, 31, 32, 33, 100):
            self.assertEqual(probe("stream", self.home, "pass-ä-1", nonce.hex(), str(length)), seal.stream(nonce, length).hex())

    def test_python_seals_cpp_opens(self):
        seal = tai.Seal(self.home, "pw-1")
        for text in ("", "a", "hello wörld \U0001F600", "x" * 5000, "line\nbreak$dollar"):
            self.assertEqual(probe("unlock", self.home, "pw-1", seal.lock(text)), text.encode().hex())

    def test_cpp_seals_python_opens(self):
        seal = tai.Seal(self.home, "pw-2")
        for text in ("", "a", "hello wörld \U0001F600", "x" * 5000):
            self.assertEqual(seal.unlock(probe("lock", self.home, "pw-2", text.encode().hex())), text)

    def test_cpp_creates_seal_file_python_opens(self):
        probe("master", self.home, "fresh")
        raw = open(os.path.join(self.home, ".seal")).read()
        self.assertEqual(list(json.loads(raw)), ["salt", "check"])
        self.assertEqual(raw, json.dumps(json.loads(raw)))
        tai.Seal(self.home, "fresh")
        with self.assertRaises(tai.SealError):
            tai.Seal(self.home, "other")
        self.assertEqual(probe("master", self.home, "other"), "ERR wrong TAI_PASSPHRASE")
        self.assertEqual(oct(os.stat(os.path.join(self.home, ".seal")).st_mode & 0o777), "0o600")

    def test_wrong_passphrase_and_tamper(self):
        seal = tai.Seal(self.home, "pw-3")
        self.assertEqual(probe("unlock", self.home, "wrong", "tai1$x"), "ERR wrong TAI_PASSPHRASE")
        locked = seal.lock("precious")
        bad = locked[:-4] + ("A" if locked[-4] != "A" else "B") + locked[-3:]
        self.assertEqual(probe("unlock", self.home, "pw-3", bad), "ERR seal check failed")
        self.assertEqual(probe("unlock", self.home, "pw-3", "plain"), "ERR not sealed")
        self.assertEqual(probe("unlock", self.home, "pw-3", "tai1$a$b"), "ERR sealed value is corrupt")
        with open(os.path.join(self.home, ".seal"), "w") as handle:
            handle.write("{broken")
        self.assertEqual(probe("master", self.home, "pw-3"), "ERR seal file is corrupt")

    def populate_python(self, passphrase):
        config = tai.Config(Args())
        store = tai.Store(config, tai.Seal(config.home, passphrase))
        store.log_event("default", "user", "message", "event with ünïcode and orchid-4421")
        store.save_secret("api", "secret-value-9", {"kind": "token"})
        store.add_schedule("job", "scheduled prompt zebra", "t", 3600, "2026-01-01T00:00:00+00:00", 60)
        store.save_system("work", "system text kiwi")
        store.save_session("work", [{"role": "user", "content": "session text"}])
        store.mcp_save_server({"name": "srv", "transport": "http", "url": "http://x", "headers": {"Authorization": "Bearer abc"}})
        store.upsert_settings({"backends": "lbl|http://b|key"})
        store.close()

    def dump_python(self, passphrase):
        seal = tai.Seal(self.home, passphrase)
        db = sqlite3.connect(os.path.join(self.home, "memory.db"))

        def col(sql):
            return [("S:" + seal.unlock(row[0])) if row[0].startswith("tai1$") else "P:" + row[0] for row in db.execute(sql).fetchall()]

        doc = {
            "events": col("SELECT text FROM events ORDER BY id"),
            "secret_value": col("SELECT value FROM secrets ORDER BY name"),
            "secret_meta": col("SELECT meta FROM secrets WHERE meta IS NOT NULL ORDER BY name"),
            "schedules": col("SELECT prompt FROM schedules ORDER BY id"),
            "settings": col("SELECT value FROM settings WHERE key = 'backends'"),
            "mcp_headers": col("SELECT headers_json FROM mcp_servers ORDER BY name"),
            "files": {},
        }
        db.close()
        pdir = os.path.join(self.home, "profiles")
        for name in sorted(os.listdir(pdir)):
            raw = open(os.path.join(pdir, name), encoding="utf-8").read()
            doc["files"][name] = ("S:" + seal.unlock(raw)) if raw.startswith("tai1$") else "P:" + raw
        return doc

    def dump_cpp(self, passphrase):
        return json.loads(probe("vault_dump", self.home, passphrase))

    def test_vault_python_written_cpp_read(self):
        self.populate_python("vault-1")
        py = self.dump_python("vault-1")
        self.assertTrue(all(v.startswith("S:") for v in py["events"] + py["secret_value"] + py["secret_meta"] + py["schedules"] + py["settings"] + py["mcp_headers"]))
        self.assertEqual(self.dump_cpp("vault-1"), py)

    def test_vault_cpp_written_python_read(self):
        self.assertEqual(probe("vault_write", self.home, "vault-2"), "ok")
        cpp = self.dump_cpp("vault-2")
        self.assertTrue(cpp["events"][0].startswith("S:"))
        self.assertEqual(self.dump_python("vault-2"), cpp)
        config = tai.Config(Args())
        store = tai.Store(config, tai.Seal(config.home, "vault-2"))
        self.assertEqual(store.load_secret("api"), "secret-value-9")
        self.assertEqual(store.load_system("work"), "system text kiwi")
        self.assertEqual(store.load_session("work"), [{"role": "user", "content": "session text"}])
        store.close()

    def test_rotate_cpp_then_python_reads(self):
        self.populate_python(tai.DEFAULT_PASSPHRASE)
        before = self.dump_python(tai.DEFAULT_PASSPHRASE)
        self.assertEqual(probe("rotate", self.home, "personal-9"), "ok")
        with self.assertRaises(tai.SealError):
            tai.Seal(self.home, tai.DEFAULT_PASSPHRASE)
        self.assertEqual(self.dump_python("personal-9"), before)

    def test_rotate_python_then_cpp_reads(self):
        self.populate_python(tai.DEFAULT_PASSPHRASE)
        before = self.dump_python(tai.DEFAULT_PASSPHRASE)
        config = tai.Config(Args())
        tai.rotate_seal(config, tai.Seal(config.home, tai.DEFAULT_PASSPHRASE), "personal-8")
        self.assertEqual(self.dump_cpp("personal-8"), before)

    def test_boot_upgrade_factory(self):
        self.populate_python(tai.DEFAULT_PASSPHRASE)
        before = self.dump_python(tai.DEFAULT_PASSPHRASE)
        out = subprocess.run([PROBE, "factory", self.home, "boot-pass", "0"], capture_output=True, text=True).stdout
        self.assertIn("seal upgraded from default key to TAI_PASSPHRASE", out)
        self.assertIn("ok enabled=1 default=0", out)
        self.assertEqual(self.dump_python("boot-pass"), before)
        self.assertEqual(probe("factory", self.home, "wrong", "0"), "ERR wrong TAI_PASSPHRASE")
        self.assertEqual(probe("factory", self.home, tai.DEFAULT_PASSPHRASE, "1"), "ERR wrong TAI_PASSPHRASE")
        self.assertEqual(probe("factory", self.home, "", "0"), "ok enabled=0 default=0")


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