from __future__ import annotations from dataclasses import dataclass, replace from datetime import UTC, datetime import hashlib from types import SimpleNamespace import unittest from sqlalchemy import create_engine from sqlalchemy.orm import sessionmaker from govoplan_core.core.voting import ( VotingBallotCreateCommand, VotingCastCommand, VotingElector, VotingOption, voting_provider_capability, ) from govoplan_core.core.encryption import ( CAPABILITY_ENCRYPTION_CONTENT_CIPHER, CAPABILITY_ENCRYPTION_KEY_VAULT, ContentProtectionEnvelope, ProtectedContent, ) from govoplan_core.db.base import Base from govoplan_voting.backend.db.models import ( VotingCastRecord, VotingConfidentialCast, ) from govoplan_voting.backend.local_confidential_provider import ( LOCAL_CONFIDENTIAL_PROVIDER_ID, LocalConfidentialVotingProvider, ) from govoplan_voting.backend.service import SqlVotingBallots, VotingStoreError @dataclass class Principal: tenant_id: str account_id: str class FakeRegistry: def __init__(self) -> None: self.capabilities: dict[str, object] = {} def has_capability(self, name: str) -> bool: return name in self.capabilities def capability(self, name: str) -> object | None: return self.capabilities.get(name) class FakeKeyVault: def __init__(self) -> None: self.vaults: dict[tuple[str, str], object] = {} def create_vault(self, session, principal, *, request): del session, principal vault = SimpleNamespace(state="active", current_key=SimpleNamespace(version=1)) self.vaults[(request.tenant_id, request.vault_id)] = vault return vault def get_vault(self, session, *, tenant_id, vault_id): del session return self.vaults.get((tenant_id, vault_id)) def rotate_key(self, *args, **kwargs): raise NotImplementedError def revoke_key(self, *args, **kwargs): raise NotImplementedError def schedule_destruction(self, *args, **kwargs): raise NotImplementedError def reconcile_vault(self, *args, **kwargs): raise NotImplementedError class FakeContentCipher: def __init__(self) -> None: self.plaintexts: dict[str, bytes] = {} def protect_content(self, session, *, request): del session envelope_id = f"envelope-{request.resource_id}" ciphertext = b"ciphertext:" + hashlib.sha256(request.plaintext).digest() self.plaintexts[envelope_id] = request.plaintext return ProtectedContent( envelope=ContentProtectionEnvelope( envelope_id=envelope_id, tenant_id=request.tenant_id, owner_module=request.owner_module, resource_type=request.resource_type, resource_id=request.resource_id, profile_kind="server_envelope", profile_id=request.profile_id, provider_id="fake", vault_id=request.vault_id, key_version=1, algorithm_suite="AES-256-GCM", ciphertext_ref=request.ciphertext_ref, ciphertext_digest=hashlib.sha256(ciphertext).hexdigest(), authenticated_context_digest="a" * 64, state="active", created_at=datetime.now(UTC), wrapped_key_refs=(f"wrapped-{request.resource_id}",), ), ciphertext=ciphertext, ) def unprotect_content(self, session, *, request): del session return self.plaintexts[request.envelope_id] def execute_rewrap(self, *args, **kwargs): raise NotImplementedError def prepare_reencryption(self, *args, **kwargs): raise NotImplementedError def command(*, assurance: str = "recorded", provider_id: str | None = None): return VotingBallotCreateCommand( title="Budget choice", method="single_choice", assurance_profile=assurance, options=(VotingOption("a", "Option A"), VotingOption("b", "Option B")), electorate=( VotingElector("alice", "Alice", 2, {"source": "mandate:1"}), VotingElector("bob", "Bob", 1, {"source": "mandate:2"}), ), quorum_weight=2, threshold_numerator=1, threshold_denominator=2, provider_id=provider_id, provider_ballot_ref="provider-ballot-1" if provider_id else None, ) class VotingTests(unittest.TestCase): def setUp(self) -> None: self.engine = create_engine("sqlite+pysqlite:///:memory:") Base.metadata.create_all(self.engine) self.Session = sessionmaker(bind=self.engine) self.service = SqlVotingBallots() self.manager = Principal("tenant-1", "manager") def tearDown(self) -> None: Base.metadata.drop_all(self.engine) self.engine.dispose() def create_open(self, session): created = self.service.create_ballot( session, self.manager, command=command(), idempotency_key="create-1", ) opened = self.service.open_ballot( session, self.manager, ballot_id=created.id, expected_revision=created.revision, idempotency_key="open-1", ) return opened def test_recorded_ballot_freezes_replaces_tallies_and_certifies(self) -> None: with self.Session() as session: opened = self.create_open(session) detail = self.service.get_ballot(session, self.manager, ballot_id=opened.id) self.assertEqual(64, len(str(detail["definition_sha256"]))) self.assertEqual(64, len(str(detail["electorate_sha256"]))) first = self.service.cast_ballot( session, Principal("tenant-1", "alice"), ballot_id=opened.id, command=VotingCastCommand(("a",), "cast-alice-1"), ) replay = self.service.cast_ballot( session, Principal("tenant-1", "alice"), ballot_id=opened.id, command=VotingCastCommand(("a",), "cast-alice-1"), ) self.assertEqual(first.receipt_sha256, replay.receipt_sha256) self.assertTrue(replay.replayed) replaced = self.service.cast_ballot( session, Principal("tenant-1", "alice"), ballot_id=opened.id, command=VotingCastCommand(("b",), "cast-alice-2"), ) self.assertTrue(replaced.replaced_previous) self.service.cast_ballot( session, Principal("tenant-1", "bob"), ballot_id=opened.id, command=VotingCastCommand(("b",), "cast-bob-1"), ) result = self.service.close_ballot( session, self.manager, ballot_id=opened.id, expected_revision=opened.revision, idempotency_key="close-1", ) self.assertEqual({"a": 0, "b": 2}, dict(result.counts)) self.assertEqual({"a": 0, "b": 3}, dict(result.weighted_counts)) self.assertTrue(result.quorum_met) self.assertTrue(result.threshold_met) certified = self.service.certify_ballot( session, self.manager, ballot_id=opened.id, expected_revision=result.revision, evidence=(), idempotency_key="certify-1", ) self.assertEqual("certified", certified.state) events = self.service.history(session, self.manager, ballot_id=opened.id) self.assertEqual( [ "ballot.created", "ballot.opened", "ballot.vote_cast", "ballot.vote_cast", "ballot.vote_cast", "ballot.closed", "ballot.certified", ], [item["event_type"] for item in events], ) active = ( session.query(VotingCastRecord) .filter( VotingCastRecord.ballot_id == opened.id, VotingCastRecord.superseded_at.is_(None), ) .count() ) self.assertEqual(2, active) def test_eligibility_tenant_and_revision_are_enforced(self) -> None: with self.Session() as session: opened = self.create_open(session) with self.assertRaisesRegex(VotingStoreError, "frozen electorate"): self.service.cast_ballot( session, Principal("tenant-1", "mallory"), ballot_id=opened.id, command=VotingCastCommand(("a",), "cast-mallory"), ) self.assertIsNone( self.service.get_ballot( session, Principal("tenant-2", "manager"), ballot_id=opened.id, ) ) with self.assertRaisesRegex(VotingStoreError, "revision conflict"): self.service.close_ballot( session, self.manager, ballot_id=opened.id, expected_revision=1, idempotency_key="close-stale", ) def test_secret_profile_fails_closed_without_provider(self) -> None: with self.Session() as session: created = self.service.create_ballot( session, self.manager, command=command(assurance="secret", provider_id="certified"), idempotency_key="create-secret", ) with self.assertRaisesRegex( VotingStoreError, "requires an available external provider" ): self.service.open_ballot( session, self.manager, ballot_id=created.id, expected_revision=created.revision, idempotency_key="open-secret", ) def test_local_confidential_provider_encrypts_casts_and_returns_aggregates( self, ) -> None: registry = FakeRegistry() provider = LocalConfidentialVotingProvider(registry) registry.capabilities.update( { voting_provider_capability(LOCAL_CONFIDENTIAL_PROVIDER_ID): provider, CAPABILITY_ENCRYPTION_KEY_VAULT: FakeKeyVault(), CAPABILITY_ENCRYPTION_CONTENT_CIPHER: FakeContentCipher(), } ) service = SqlVotingBallots(registry) draft = replace( command( assurance="confidential", provider_id=LOCAL_CONFIDENTIAL_PROVIDER_ID, ), provider_ballot_ref=None, ) with self.Session() as session: created = service.create_ballot( session, self.manager, command=draft, idempotency_key="create-confidential", ) opened = service.open_ballot( session, self.manager, ballot_id=created.id, expected_revision=created.revision, idempotency_key="open-confidential", ) detail = service.get_ballot( session, self.manager, ballot_id=opened.id, ) self.assertEqual( f"local-confidential:{opened.id}", detail["provider_ballot_ref"], ) first = service.cast_ballot( session, Principal("tenant-1", "alice"), ballot_id=opened.id, command=VotingCastCommand(("a",), "cast-confidential-alice"), ) replay = service.cast_ballot( session, Principal("tenant-1", "alice"), ballot_id=opened.id, command=VotingCastCommand(("a",), "cast-confidential-alice"), ) self.assertEqual(first.receipt_sha256, replay.receipt_sha256) self.assertTrue(replay.replayed) with self.assertRaisesRegex( VotingStoreError, "provider cast was rejected", ): service.cast_ballot( session, Principal("tenant-1", "alice"), ballot_id=opened.id, command=VotingCastCommand( ("b",), "cast-confidential-alice", ), ) service.cast_ballot( session, Principal("tenant-1", "bob"), ballot_id=opened.id, command=VotingCastCommand(("b",), "cast-confidential-bob"), ) self.assertEqual(0, session.query(VotingCastRecord).count()) encrypted = session.query(VotingConfidentialCast).all() self.assertEqual(2, len(encrypted)) self.assertTrue(all(row.ciphertext for row in encrypted)) result = service.close_ballot( session, self.manager, ballot_id=opened.id, expected_revision=opened.revision, idempotency_key="close-confidential", ) self.assertEqual({"a": 1, "b": 1}, dict(result.counts)) self.assertEqual({"a": 2, "b": 1}, dict(result.weighted_counts)) self.assertEqual("local_confidential", result.evidence[0]["provider_id"]) self.assertFalse(result.evidence[0]["certified"]) def test_idempotency_key_cannot_change_command(self) -> None: with self.Session() as session: self.service.create_ballot( session, self.manager, command=command(), idempotency_key="create-replay", ) changed = replace(command(), title="Different") with self.assertRaisesRegex(VotingStoreError, "idempotency key"): self.service.create_ballot( session, self.manager, command=changed, idempotency_key="create-replay", ) if __name__ == "__main__": unittest.main()