tai/test_seal.py

# retoor <retoor@molodetz.nl>
import os
import sqlite3
import sys
import tempfile
import unittest

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import tai


class FakeArgs:
    profile = "default"
    yes = True


class SealTests(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        os.environ["TAI_HOME"] = self.tmp.name

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

    def make_store(self, passphrase="test-secret-1"):
        config = tai.Config(FakeArgs())
        seal = tai.Seal(config.home, passphrase)
        return tai.Store(config, seal)

    def test_round_trip(self):
        config = tai.Config(FakeArgs())
        seal = tai.Seal(config.home, "test-secret-1")
        locked = seal.lock("the wifi password is orchid-4421")
        self.assertTrue(locked.startswith("tai1$"))
        self.assertNotIn("orchid", locked)
        self.assertEqual(seal.unlock(locked), "the wifi password is orchid-4421")

    def test_unique_nonce(self):
        config = tai.Config(FakeArgs())
        seal = tai.Seal(config.home, "test-secret-1")
        self.assertNotEqual(seal.lock("same"), seal.lock("same"))

    def test_wrong_passphrase(self):
        config = tai.Config(FakeArgs())
        tai.Seal(config.home, "correct-horse")
        with self.assertRaises(tai.SealError):
            tai.Seal(config.home, "wrong-guess")

    def test_tamper_detected(self):
        config = tai.Config(FakeArgs())
        seal = tai.Seal(config.home, "test-secret-1")
        locked = seal.lock("precious")
        tampered = locked[:-4] + ("A" if locked[-4] != "A" else "B") + locked[-3:]
        with self.assertRaises(tai.SealError):
            seal.unlock(tampered)

    def test_sql_functions(self):
        store = self.make_store()
        row = store.db.execute("SELECT tai_enc('hello-sql')").fetchone()[0]
        self.assertTrue(row.startswith("tai1$"))
        plain = store.db.execute("SELECT tai_dec(?)", (row,)).fetchone()[0]
        self.assertEqual(plain, "hello-sql")
        store.close()

    def test_sealed_events_recall(self):
        store = self.make_store()
        store.log_event("default", "user", "message", "my vault code is zebra-9")
        raw = store.db.execute("SELECT text FROM events").fetchone()[0]
        self.assertTrue(raw.startswith("tai1$"))
        self.assertNotIn("zebra", raw)
        rows = store.search_events("default", "zebra")
        self.assertEqual(len(rows), 1)
        self.assertIn("zebra-9", rows[0][3])
        store.close()

    def test_sealed_profile_files(self):
        store = self.make_store()
        store.save_system("work", "system knows orchid-4421")
        with open(store.profile_path("work"), encoding="utf-8") as handle:
            self.assertTrue(handle.read().startswith("tai1$"))
        self.assertEqual(store.load_system("work"), "system knows orchid-4421")
        store.save_session("work", [{"role": "user", "content": "secret hello"}])
        with open(store.session_path("work"), encoding="utf-8") as handle:
            self.assertTrue(handle.read().startswith("tai1$"))
        self.assertEqual(store.load_session("work"), [{"role": "user", "content": "secret hello"}])
        store.close()

    def test_migration_from_plaintext(self):
        plain = self.make_store(passphrase="")
        plain.log_event("default", "user", "message", "plain memory")
        plain.save_system("default", "plain system")
        plain.close()
        sealed = self.make_store(passphrase="new-secret-2")
        raw = sealed.db.execute("SELECT text FROM events").fetchone()[0]
        self.assertTrue(raw.startswith("tai1$"))
        with open(sealed.profile_path("default"), encoding="utf-8") as handle:
            self.assertTrue(handle.read().startswith("tai1$"))
        self.assertEqual(sealed.load_system("default"), "plain system")
        sealed.close()

    def test_sealed_home_requires_passphrase(self):
        sealed = self.make_store()
        sealed.log_event("default", "user", "message", "locked memory")
        sealed.close()
        config = tai.Config(FakeArgs())
        with self.assertRaises(tai.SealError):
            tai.Store(config, tai.Seal(config.home, ""))

    def test_passphrase_default(self):
        os.environ.pop("TAI_PASSPHRASE", None)
        value, default = tai.resolve_passphrase()
        self.assertEqual(value, tai.DEFAULT_PASSPHRASE)
        self.assertTrue(default)
        os.environ["TAI_PASSPHRASE"] = "personal-1"
        value, default = tai.resolve_passphrase()
        self.assertEqual(value, "personal-1")
        self.assertFalse(default)
        os.environ["TAI_PASSPHRASE"] = ""
        value, default = tai.resolve_passphrase()
        self.assertEqual(value, "")
        self.assertFalse(default)
        os.environ.pop("TAI_PASSPHRASE", None)

    def test_rotate_default_to_personal(self):
        os.environ.pop("TAI_PASSPHRASE", None)
        config = tai.Config(FakeArgs())
        old = tai.Seal(config.home, tai.DEFAULT_PASSPHRASE)
        store = tai.Store(config, old)
        store.log_event("default", "user", "message", "rotate me")
        store.save_system("default", "rotate system")
        store.close()
        fresh = tai.rotate_seal(config, old, "personal-2")
        self.assertTrue(fresh.enabled)
        with self.assertRaises(tai.SealError):
            tai.Seal(config.home, tai.DEFAULT_PASSPHRASE)
        reopened = tai.Store(config, tai.Seal(config.home, "personal-2"))
        self.assertEqual(reopened.load_system("default"), "rotate system")
        self.assertIn("rotate me", reopened.search_events("default", "rotate")[0][3])
        raw = reopened.db.execute("SELECT text FROM events").fetchone()[0]
        self.assertTrue(raw.startswith("tai1$"))
        reopened.close()

    def test_plaintext_home_unaffected(self):
        store = self.make_store(passphrase="")
        store.log_event("default", "user", "message", "open memory")
        self.assertEqual(store.search_events("default", "open")[0][3], "open memory")
        store.save_system("default", "open system")
        with open(store.profile_path("default"), encoding="utf-8") as handle:
            self.assertEqual(handle.read(), "open system")
        store.close()


    def test_rotate_moves_secrets(self):
        os.environ.pop("TAI_PASSPHRASE", None)
        config = tai.Config(FakeArgs())
        old = tai.Seal(config.home, tai.DEFAULT_PASSPHRASE)
        store = tai.Store(config, old)
        store.save_secret("api", "rotate-me-1")
        store.close()
        fresh = tai.rotate_seal(config, old, "personal-3")
        reopened = tai.Store(config, tai.Seal(config.home, "personal-3"))
        self.assertEqual(reopened.load_secret("api"), "rotate-me-1")
        raw = reopened.db.execute("SELECT value FROM secrets WHERE name = 'api'").fetchone()[0]
        self.assertTrue(raw.startswith("tai1$"))
        reopened.close()

    def test_sealed_secrets_require_passphrase(self):
        sealed = self.make_store()
        sealed.save_secret("k", "locked-value")
        sealed.close()
        config = tai.Config(FakeArgs())
        with self.assertRaises(tai.SealError):
            tai.Store(config, tai.Seal(config.home, ""))

    def test_schedule_prompt_sealed(self):
        store = self.make_store()
        store.add_schedule("job", "prompt-secret-3", "t", 0, "2026-01-01T00:00:00+00:00", 60)
        raw = store.db.execute("SELECT prompt FROM schedules").fetchone()[0]
        self.assertTrue(raw.startswith("tai1$"))
        self.assertNotIn("prompt-secret", raw)
        self.assertEqual(store.list_schedules("t")[0]["prompt"], "prompt-secret-3")
        store.close()

    def test_rotate_moves_schedules(self):
        os.environ.pop("TAI_PASSPHRASE", None)
        config = tai.Config(FakeArgs())
        old = tai.Seal(config.home, tai.DEFAULT_PASSPHRASE)
        store = tai.Store(config, old)
        store.add_schedule("job", "rotate-prompt-4", "t", 3600, "2026-01-01T00:00:00+00:00", 60)
        store.close()
        fresh = tai.rotate_seal(config, old, "personal-4")
        reopened = tai.Store(config, tai.Seal(config.home, "personal-4"))
        self.assertEqual(reopened.list_schedules("t")[0]["prompt"], "rotate-prompt-4")
        raw = reopened.db.execute("SELECT prompt FROM schedules").fetchone()[0]
        self.assertTrue(raw.startswith("tai1$"))
        reopened.close()


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