from __future__ import annotations from datetime import UTC, datetime, timedelta import hashlib import unittest from sqlalchemy import create_engine, select from sqlalchemy.orm import Session from govoplan_core.core.encryption import ( CAPABILITY_ENCRYPTION_CONTENT_CIPHER_PROVIDER_PREFIX, CAPABILITY_ENCRYPTION_KEY_MATERIAL_PROVIDER_PREFIX, ContentProtectionRequest, ContentUnprotectionRequest, KeyRotationRequest, KeyVaultCreateRequest, ProtectionMigrationRequest, ) from govoplan_core.core.identity_trust import ( CAPABILITY_IDENTITY_TRUST_ASSURANCE, AssuranceDecision, ) from govoplan_encryption.backend.db.models import ( ContentProtectionRecord, EncryptionKeyOperation, EncryptionKeyVersion, EncryptionLocalKeyMaterial, EncryptionLocalProviderOperation, EncryptionLocalWrappedContentKey, EncryptionVault, ProtectionMigration, RecoveryApproval, RecoveryCeremony, ) from govoplan_encryption.backend.local_provider import ( LOCAL_PROVIDER_ID, LocalAesGcmProvider, LocalProviderError, ) from govoplan_encryption.backend.service import SqlEncryptionService class Principal: tenant_id = "tenant-1" account_id = "account-1" class Assurance: def verify_assurance(self, _session, _principal, *, request): return AssuranceDecision( allowed=True, reason="test assurance", assurance_level="high", evidence_ref=request.evidence_ref, verified_at=request.evaluated_at, expires_at=request.evaluated_at + timedelta(minutes=5), provenance={"provider_id": "test"}, ) class Registry: def __init__(self, provider: LocalAesGcmProvider) -> None: self.values = { f"{CAPABILITY_ENCRYPTION_KEY_MATERIAL_PROVIDER_PREFIX}{LOCAL_PROVIDER_ID}": provider, f"{CAPABILITY_ENCRYPTION_CONTENT_CIPHER_PROVIDER_PREFIX}{LOCAL_PROVIDER_ID}": provider, CAPABILITY_IDENTITY_TRUST_ASSURANCE: Assurance(), } def has_capability(self, name: str) -> bool: return name in self.values def capability(self, name: str): return self.values[name] class SequentialReferences: def __init__(self) -> None: self.value = 0 def __call__(self) -> str: self.value += 1 return f"reference-{self.value}" class LocalProviderTests(unittest.TestCase): def setUp(self) -> None: self.engine = create_engine("sqlite+pysqlite:///:memory:") self.tables = [ EncryptionVault.__table__, EncryptionKeyVersion.__table__, EncryptionKeyOperation.__table__, EncryptionLocalKeyMaterial.__table__, EncryptionLocalWrappedContentKey.__table__, EncryptionLocalProviderOperation.__table__, ContentProtectionRecord.__table__, ProtectionMigration.__table__, RecoveryCeremony.__table__, RecoveryApproval.__table__, ] for table in self.tables: table.create(self.engine) self.session = Session(self.engine) self.provider = LocalAesGcmProvider( bytes(range(32)), random_bytes=lambda size: bytes([size]) * size, reference_factory=SequentialReferences(), ) self.registry = Registry(self.provider) self.service = SqlEncryptionService(self.registry) self.principal = Principal() self._create_vault() self.session.commit() def tearDown(self) -> None: self.session.close() self.engine.dispose() def _create_vault(self) -> None: result = self.service.create_vault( self.session, self.principal, request=KeyVaultCreateRequest( tenant_id="tenant-1", vault_id="vault-1", name="Reference content vault", provider_id=LOCAL_PROVIDER_ID, purpose="feature-content", algorithm_suite="AES-256-GCM", scope_type="tenant", scope_id="tenant-1", policy_ref="policy:encryption:v1", idempotency_key="create-vault-1", ), ) self.assertEqual("active", result.state) def _protect(self, *, idempotency_key: str = "protect-1"): return self.service.protect_content( self.session, request=ContentProtectionRequest( tenant_id="tenant-1", owner_module="files", resource_type="file_blob", resource_id="blob-1", profile_id="server-envelope-v1", vault_id="vault-1", ciphertext_ref="object://files/blob-1", plaintext=b"sensitive monthly records", policy_decision_ref="policy:protect:v1", idempotency_key=idempotency_key, actor_id="account-1", metadata={"content_type": "text/plain"}, ), ) def _open(self, protected) -> bytes: return self.service.unprotect_content( self.session, request=ContentUnprotectionRequest( tenant_id="tenant-1", owner_module="files", resource_type="file_blob", resource_id="blob-1", envelope_id=protected.envelope.envelope_id, ciphertext=protected.ciphertext, actor_id="account-1", ), ) def test_round_trip_uses_wrapped_keys_and_rejects_tampering(self) -> None: protected = self._protect() self.assertNotIn(b"sensitive monthly records", protected.ciphertext) self.assertEqual( "1c4d240db81a06892585447218d83e0fbc53d5d346206d758351b2af635cbe01108699e93edf4057c8", protected.ciphertext.hex(), ) self.assertEqual(b"sensitive monthly records", self._open(protected)) key = self.session.scalar(select(EncryptionLocalKeyMaterial)) wrapped = self.session.scalar(select(EncryptionLocalWrappedContentKey)) self.assertIsNotNone(key.wrapped_key) self.assertIsNotNone(wrapped.wrapped_content_key) self.assertNotEqual(bytes(range(32)), key.wrapped_key) self.assertNotEqual( hashlib.sha256(b"sensitive monthly records").hexdigest(), wrapped.plaintext_commitment, ) corrupted = bytearray(protected.ciphertext) corrupted[-1] ^= 1 with self.assertRaisesRegex(LocalProviderError, "digest verification"): self.service.unprotect_content( self.session, request=ContentUnprotectionRequest( tenant_id="tenant-1", owner_module="files", resource_type="file_blob", resource_id="blob-1", envelope_id=protected.envelope.envelope_id, ciphertext=bytes(corrupted), ), ) def test_idempotent_replay_returns_identical_ciphertext(self) -> None: first = self._protect() replay = self._protect() self.assertTrue(replay.replayed) self.assertEqual(first.envelope.envelope_id, replay.envelope.envelope_id) self.assertEqual(first.ciphertext, replay.ciphertext) self.assertEqual( 1, self.session.query(EncryptionLocalWrappedContentKey).count(), ) def test_rotation_rewraps_without_changing_ciphertext(self) -> None: protected = self._protect() rotated = self.service.rotate_key( self.session, self.principal, request=KeyRotationRequest( tenant_id="tenant-1", vault_id="vault-1", expected_revision=1, reason="scheduled rotation", policy_decision_ref="policy:rotate:v1", assurance_evidence_ref="assurance:rotate:v1", idempotency_key="rotate-vault-1", ), ) self.assertEqual(2, rotated.current_key.version) migration = self.service.request_migration( self.session, self.principal, request=ProtectionMigrationRequest( tenant_id="tenant-1", envelope_id=protected.envelope.envelope_id, target_provider_id=LOCAL_PROVIDER_ID, target_vault_id="vault-1", target_key_version=2, target_algorithm_suite="AES-256-GCM", mode="rewrap", policy_decision_ref="policy:rewrap:v1", assurance_evidence_ref="assurance:rewrap:v1", idempotency_key="rewrap-1", ), ) result = self.service.execute_rewrap( self.session, self.principal, migration_id=migration.migration_id, ) self.assertEqual("succeeded", result.state) self.assertEqual( protected.envelope.ciphertext_digest, result.target_envelope.ciphertext_digest, ) opened = self.service.unprotect_content( self.session, request=ContentUnprotectionRequest( tenant_id="tenant-1", owner_module="files", resource_type="file_blob", resource_id="blob-1", envelope_id=result.target_envelope.envelope_id, ciphertext=protected.ciphertext, ), ) self.assertEqual(b"sensitive monthly records", opened) def test_reencryption_waits_for_owner_confirmation_and_replays_target(self) -> None: protected = self._protect() rotated = self.service.rotate_key( self.session, self.principal, request=KeyRotationRequest( tenant_id="tenant-1", vault_id="vault-1", expected_revision=1, reason="algorithm-preserving content rotation", policy_decision_ref="policy:rotate:v1", assurance_evidence_ref="assurance:rotate:v1", idempotency_key="rotate-vault-for-reencrypt", ), ) migration = self.service.request_migration( self.session, self.principal, request=ProtectionMigrationRequest( tenant_id="tenant-1", envelope_id=protected.envelope.envelope_id, target_provider_id=LOCAL_PROVIDER_ID, target_vault_id="vault-1", target_key_version=rotated.current_key.version, target_algorithm_suite="AES-256-GCM", mode="reencrypt", policy_decision_ref="policy:reencrypt:v1", assurance_evidence_ref="assurance:reencrypt:v1", idempotency_key="reencrypt-1", ), ) prepared = self.service.prepare_reencryption( self.session, self.principal, migration_id=migration.migration_id, source_ciphertext=protected.ciphertext, target_ciphertext_ref="object://files/blob-1-v2", ) replay = self.service.prepare_reencryption( self.session, self.principal, migration_id=migration.migration_id, source_ciphertext=protected.ciphertext, target_ciphertext_ref="object://files/blob-1-v2", ) self.assertNotEqual(protected.ciphertext, prepared.ciphertext) self.assertTrue(replay.replayed) self.assertEqual(prepared.ciphertext, replay.ciphertext) self.assertIsNone( self.service.get_envelope( self.session, self.principal, tenant_id="tenant-1", envelope_id=prepared.envelope.envelope_id, ) ) unresolved = self.service.record_migration_outcome( self.session, self.principal, migration_id=migration.migration_id, state="outcome_unknown", evidence_refs=("owner-store:write-outcome-unknown",), error_code="storage_timeout", ) self.assertEqual("outcome_unknown", unresolved.state) prepared = self.service.prepare_reencryption( self.session, self.principal, migration_id=migration.migration_id, source_ciphertext=protected.ciphertext, target_ciphertext_ref="object://files/blob-1-v2", ) self.assertTrue(prepared.replayed) completed = self.service.record_migration_outcome( self.session, self.principal, migration_id=migration.migration_id, state="succeeded", evidence_refs=("owner-store:files:blob-1-v2",), target_envelope=prepared.envelope, ) self.assertEqual("succeeded", completed.state) opened = self.service.unprotect_content( self.session, request=ContentUnprotectionRequest( tenant_id="tenant-1", owner_module="files", resource_type="file_blob", resource_id="blob-1", envelope_id=prepared.envelope.envelope_id, ciphertext=prepared.ciphertext, ), ) self.assertEqual(b"sensitive monthly records", opened) def test_database_restore_requires_the_matching_deployment_key(self) -> None: protected = self._protect() self.session.commit() restored_engine = create_engine("sqlite+pysqlite:///:memory:") source = self.engine.raw_connection() target = restored_engine.raw_connection() try: source.driver_connection.backup(target.driver_connection) target.commit() finally: target.close() source.close() restored_session = Session(restored_engine) restored_service = SqlEncryptionService( Registry(LocalAesGcmProvider(bytes(range(32)))) ) request = ContentUnprotectionRequest( tenant_id="tenant-1", owner_module="files", resource_type="file_blob", resource_id="blob-1", envelope_id=protected.envelope.envelope_id, ciphertext=protected.ciphertext, ) self.assertEqual( b"sensitive monthly records", restored_service.unprotect_content(restored_session, request=request), ) wrong_service = SqlEncryptionService( Registry(LocalAesGcmProvider(bytes(reversed(range(32))))) ) try: with self.assertRaisesRegex( LocalProviderError, "configured deployment key", ): wrong_service.unprotect_content(restored_session, request=request) finally: restored_session.close() restored_engine.dispose() def test_destroyed_vault_key_fails_closed(self) -> None: protected = self._protect() key = self.session.scalar(select(EncryptionLocalKeyMaterial)) self.provider.schedule_key_destruction( self.session, self.principal, provider_key_ref=key.provider_key_ref, effective_at=datetime.now(tz=UTC) - timedelta(seconds=1), idempotency_key="destroy-vault-material-1", ) with self.assertRaisesRegex(LocalProviderError, "unavailable"): self._open(protected) def test_injected_encrypt_failure_rolls_back_without_mixed_state(self) -> None: def fail(point: str) -> None: if point == "encrypt.after_wrap": raise RuntimeError("injected provider failure") failed_service = SqlEncryptionService( Registry( LocalAesGcmProvider( bytes(range(32)), fault_injector=fail, ) ) ) with self.assertRaisesRegex(RuntimeError, "injected provider failure"): failed_service.protect_content( self.session, request=ContentProtectionRequest( tenant_id="tenant-1", owner_module="files", resource_type="file_blob", resource_id="blob-failed", profile_id="server-envelope-v1", vault_id="vault-1", ciphertext_ref="object://files/blob-failed", plaintext=b"must not become mixed state", policy_decision_ref="policy:protect:v1", idempotency_key="protect-failed", actor_id="account-1", ), ) self.session.rollback() self.assertEqual( 0, self.session.query(EncryptionLocalWrappedContentKey) .filter( EncryptionLocalWrappedContentKey.idempotency_key == "protect:protect-failed" ) .count(), ) self.assertEqual( 0, self.session.query(ContentProtectionRecord) .filter(ContentProtectionRecord.resource_id == "blob-failed") .count(), ) if __name__ == "__main__": unittest.main()