tai/tai_pkg/seal.py

"""Vault encryption seal and passphrase management."""

import base64
import hashlib
import hmac
import json
import os
import sqlite3
import struct

from tai_pkg.consts import Consts

__all__ = ["SealError", "Seal", "SealManager"]


class SealError(Exception):
    pass


class Seal:
    PREFIX = "tai1$"
    ROUNDS = 200000

    def __init__(self, home, passphrase):
        self.enabled = bool(passphrase)
        self.default_key = False
        self.master = b""
        self.enc_key = b""
        self.mac_key = b""
        if self.enabled:
            self.master = self.load_master(home, passphrase)
            self.enc_key = hmac.new(self.master, b"tai-enc-1", hashlib.sha256).digest()
            self.mac_key = hmac.new(self.master, b"tai-mac-1", hashlib.sha256).digest()

    def load_master(self, home, passphrase):
        path = os.path.join(home, ".seal")
        raw = ""
        if os.path.exists(path):
            with open(path, "r", encoding="utf-8") as handle:
                raw = handle.read()
        saved = {}
        if raw:
            try:
                saved = json.loads(raw)
                salt = base64.b64decode(saved["salt"])
            except (ValueError, KeyError):
                raise SealError("seal file is corrupt")
        else:
            salt = os.urandom(16)
        master = hashlib.pbkdf2_hmac("sha256", passphrase.encode("utf-8"), salt, self.ROUNDS, 32)
        check = hmac.new(master, b"tai-check-1", hashlib.sha256).hexdigest()
        if raw:
            if not hmac.compare_digest(saved.get("check", ""), check):
                raise SealError("wrong TAI_PASSPHRASE")
            return master
        with open(path, "w", encoding="utf-8") as handle:
            handle.write(json.dumps({"salt": base64.b64encode(salt).decode("ascii"), "check": check}))
        try:
            os.chmod(path, 0o600)
        except OSError:
            pass
        return master

    def stream(self, nonce, length):
        out = bytearray()
        counter = 0
        while len(out) < length:
            out += hmac.new(self.enc_key, nonce + struct.pack(">I", counter), hashlib.sha256).digest()
            counter += 1
        return bytes(out[:length])

    def lock(self, text):
        if not self.enabled or text.startswith(self.PREFIX):
            return text
        nonce = os.urandom(16)
        body = text.encode("utf-8")
        keystream = self.stream(nonce, len(body))
        cipher = bytes(piece ^ keystream[pos] for pos, piece in enumerate(body))
        mac = hmac.new(self.mac_key, nonce + cipher, hashlib.sha256).digest()
        parts = [base64.b64encode(blob).decode("ascii") for blob in (nonce, cipher, mac)]
        return self.PREFIX + "$".join(parts)

    def unlock(self, text):
        if not text.startswith(self.PREFIX):
            raise SealError("not sealed")
        try:
            _tag, nonce64, cipher64, mac64 = text.split("$", 3)
            nonce = base64.b64decode(nonce64)
            cipher = base64.b64decode(cipher64)
            mac = base64.b64decode(mac64)
        except ValueError:
            raise SealError("sealed value is corrupt")
        expect = hmac.new(self.mac_key, nonce + cipher, hashlib.sha256).digest()
        if not hmac.compare_digest(mac, expect):
            raise SealError("seal check failed")
        keystream = self.stream(nonce, len(cipher))
        plain = bytes(piece ^ keystream[pos] for pos, piece in enumerate(cipher))
        return plain.decode("utf-8")

    def safe_unlock(self, text):
        try:
            return self.unlock(text)
        except SealError:
            return "[sealed: decrypt failed]"


class SealManager:
    """SealManager."""

    @staticmethod
    def resolve_passphrase():
        raw = os.environ.get("TAI_PASSPHRASE")
        if raw is None:
            return Consts.DEFAULT_PASSPHRASE, True
        return raw, False

    @staticmethod
    def rotate_seal(config, old_seal, new_passphrase):
        db = sqlite3.connect(config.db_path)
        rows = db.execute("SELECT id, text FROM events").fetchall()
        tables = [row[0] for row in db.execute("SELECT name FROM sqlite_master WHERE type = 'table'").fetchall()]
        secret_cols = [row[1] for row in db.execute("PRAGMA table_info(secrets)").fetchall()] if "secrets" in tables else []
        has_profile = "profile" in secret_cols
        secret_rows = db.execute("SELECT profile, name, value FROM secrets").fetchall() if has_profile else [(None, name, value) for name, value in db.execute("SELECT name, value FROM secrets").fetchall()] if "secrets" in tables else []
        meta_rows = db.execute("SELECT profile, name, meta FROM secrets WHERE meta IS NOT NULL").fetchall() if has_profile and "meta" in secret_cols else [(None, name, meta) for name, meta in db.execute("SELECT name, meta FROM secrets WHERE meta IS NOT NULL").fetchall()] if "meta" in secret_cols else []
        planned_rows = db.execute("SELECT id, prompt FROM schedules").fetchall() if "schedules" in tables else []
        db.close()
        texts = {}
        for row_id, text in rows:
            texts[row_id] = old_seal.unlock(text) if text.startswith(Seal.PREFIX) else text
        secret_texts = {}
        for profile, name, value in secret_rows:
            secret_texts[(profile, name)] = old_seal.unlock(value) if value.startswith(Seal.PREFIX) else value
        secret_metas = {}
        for profile, name, meta in meta_rows:
            secret_metas[(profile, name)] = old_seal.unlock(meta) if meta.startswith(Seal.PREFIX) else meta
        planned_texts = {}
        for row_id, prompt in planned_rows:
            planned_texts[row_id] = old_seal.unlock(prompt) if prompt.startswith(Seal.PREFIX) else prompt
        files = {}
        for entry in os.listdir(config.profiles_dir):
            if not entry.endswith((".sys.md", ".session.json", ".bots.json")):
                continue
            path = os.path.join(config.profiles_dir, entry)
            with open(path, "r", encoding="utf-8") as handle:
                raw = handle.read()
            files[path] = old_seal.unlock(raw) if raw.startswith(Seal.PREFIX) else raw
        try:
            os.remove(os.path.join(config.home, ".seal"))
        except OSError:
            pass
        fresh = Seal(config.home, new_passphrase)
        db = sqlite3.connect(config.db_path)
        for row_id, plain in texts.items():
            db.execute("UPDATE events SET text = ? WHERE id = ?", (fresh.lock(plain), row_id))
        for (profile, name), plain in secret_texts.items():
            if profile is None:
                db.execute("UPDATE secrets SET value = ? WHERE name = ?", (fresh.lock(plain), name))
            else:
                db.execute("UPDATE secrets SET value = ? WHERE profile = ? AND name = ?", (fresh.lock(plain), profile, name))
        for (profile, name), plain in secret_metas.items():
            if profile is None:
                db.execute("UPDATE secrets SET meta = ? WHERE name = ?", (fresh.lock(plain), name))
            else:
                db.execute("UPDATE secrets SET meta = ? WHERE profile = ? AND name = ?", (fresh.lock(plain), profile, name))
        for row_id, plain in planned_texts.items():
            db.execute("UPDATE schedules SET prompt = ? WHERE id = ?", (fresh.lock(plain), row_id))
        db.commit()
        db.close()
        for path, plain in files.items():
            with open(path, "w", encoding="utf-8") as handle:
                handle.write(fresh.lock(plain))
        return fresh