Initial commit: single-file autonomous agent with sealed storage, sandbox, and telegram bot
This commit is contained in:
+156
@@ -0,0 +1,156 @@
|
||||
# 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()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user