f6b876fbe7
Co-Authored-By: Claude <noreply@anthropic.com> Co-Authored-By: XNLLLLH <XNLLLLH@users.noreply.github.com>
101 lines
3 KiB
Python
101 lines
3 KiB
Python
"""Tests for migrate_crypto_recover_prior_key (prior AES key re-encrypt + swap)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import secrets
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
from cryptography.exceptions import InvalidTag
|
|
|
|
from iai_mcp.migrate import migrate_crypto_recover_prior_key
|
|
from iai_mcp.store import MemoryStore
|
|
from iai_mcp.types import MemoryRecord, SCHEMA_VERSION_CURRENT
|
|
|
|
|
|
def _minimal_record(literal: str) -> MemoryRecord:
|
|
rid = uuid4()
|
|
now = datetime.now(timezone.utc)
|
|
return MemoryRecord(
|
|
id=rid,
|
|
tier="episodic",
|
|
literal_surface=literal,
|
|
aaak_index="",
|
|
embedding=[0.01] * 384,
|
|
structure_hv=b"\x00" * 1250,
|
|
community_id=None,
|
|
centrality=0.0,
|
|
detail_level=1,
|
|
pinned=False,
|
|
stability=0.0,
|
|
difficulty=0.0,
|
|
last_reviewed=None,
|
|
never_decay=False,
|
|
never_merge=False,
|
|
provenance=[],
|
|
created_at=now,
|
|
updated_at=now,
|
|
tags=[],
|
|
language="en",
|
|
s5_trust_score=0.5,
|
|
profile_modulation_gain={},
|
|
schema_version=SCHEMA_VERSION_CURRENT,
|
|
)
|
|
|
|
|
|
def test_recover_prior_key_atomic_swap_and_idempotent(tmp_path: Path) -> None:
|
|
root = tmp_path / "store"
|
|
root.mkdir()
|
|
key_a = secrets.token_bytes(32)
|
|
key_b = secrets.token_bytes(32)
|
|
kpath = root / ".crypto.key"
|
|
kpath.write_bytes(key_a)
|
|
os.chmod(kpath, 0o600)
|
|
|
|
store_a = MemoryStore(path=root, user_id="default")
|
|
rec = _minimal_record("verbatim-prior-key-recover")
|
|
store_a.insert(rec)
|
|
rid = rec.id
|
|
del store_a
|
|
|
|
kpath.write_bytes(key_b)
|
|
os.chmod(kpath, 0o600)
|
|
store_b = MemoryStore(path=root, user_id="default")
|
|
with pytest.raises(InvalidTag):
|
|
store_b.get(rid)
|
|
|
|
out = migrate_crypto_recover_prior_key(store_b, key_a, dry_run=False)
|
|
assert out.get("no_op") is False
|
|
assert out.get("records_staged") == 1
|
|
assert out.get("rows_needed_prior_key") == 1
|
|
|
|
got = store_b.get(rid)
|
|
assert got is not None
|
|
assert got.literal_surface == "verbatim-prior-key-recover"
|
|
|
|
out2 = migrate_crypto_recover_prior_key(store_b, key_a, dry_run=False)
|
|
assert out2.get("no_op") is True
|
|
assert out2.get("reason") == "all_rows_decrypt_with_current_key"
|
|
|
|
|
|
def test_recover_prior_key_dry_run_counts(tmp_path: Path) -> None:
|
|
root = tmp_path / "store2"
|
|
root.mkdir()
|
|
key_a = secrets.token_bytes(32)
|
|
key_b = secrets.token_bytes(32)
|
|
kpath = root / ".crypto.key"
|
|
kpath.write_bytes(key_a)
|
|
os.chmod(kpath, 0o600)
|
|
store_a = MemoryStore(path=root, user_id="default")
|
|
store_a.insert(_minimal_record("dry-run-count"))
|
|
del store_a
|
|
kpath.write_bytes(key_b)
|
|
os.chmod(kpath, 0o600)
|
|
store_b = MemoryStore(path=root, user_id="default")
|
|
out = migrate_crypto_recover_prior_key(store_b, key_a, dry_run=True)
|
|
assert out.get("dry_run") is True
|
|
assert out.get("would_stage") == 1
|
|
assert out.get("rows_needing_prior_key") == 1
|