Exemple #1
0
class ElectionGuardTestClient:
    def __init__(self, manifest_file: str, manifest_path: str, guardians: int,
                 quorom: int):
        self.NUMBER_OF_GUARDIANS = guardians
        self.QUORUM = quorom
        self.manifest_path = manifest_path
        self.manifest_file = manifest_file
        self.guardians: List[Guardian] = []
        self.key_ceremony_mediator = KeyCeremonyMediator(
            CeremonyDetails(self.NUMBER_OF_GUARDIANS, self.QUORUM))

    def setupGuardians(self):
        # Setup Guardians
        for i in range(self.NUMBER_OF_GUARDIANS):
            sequence = i + 2
            self.guardians.append(
                Guardian(
                    "guardian_" + str(sequence),
                    sequence,
                    self.NUMBER_OF_GUARDIANS,
                    self.QUORUM,
                ))

        # Attendance (Public Key Share)
        for guardian in self.guardians:
            self.key_ceremony_mediator.announce(guardian)

    def setupKeyCeromony(self):
        self.setupGuardians()
        self.key_ceremony_mediator.orchestrate(identity_auxiliary_encrypt)
        self.verification_results = self.key_ceremony_mediator.verify(
            identity_auxiliary_decrypt)
        self.joint_public_key = self.key_ceremony_mediator.publish_joint_key()

    def setupElectionBuilder(self):
        # Create an election builder instance, and configure it for a single public-private keypair.
        # in a real election, you would configure this for a group of guardians.  See Key Ceremony for more information.
        with open(os.path.join(self.manifest_path, self.manifest_file),
                  "r") as manifest:
            string_representation = manifest.read()
            election_description = ElectionDescription.from_json(
                string_representation)

        builder = ElectionBuilder(
            number_of_guardians=self.
            NUMBER_OF_GUARDIANS,  # since we will generate a single public-private keypair, we set this to 1
            quorum=self.
            QUORUM,  # since we will generate a single public-private keypair, we set this to 1
            description=election_description)

        builder.set_public_key(self.joint_public_key)
        self.metadata, self.context = get_optional(builder.build())
        self.builder = builder
Exemple #2
0
    def get_hamilton_election_with_encryption_context(
        self, ) -> Tuple[AllPublicElectionData, AllPrivateElectionData]:
        guardians: List[Guardian] = []
        coefficient_validation_sets: List[CoefficientValidationSet] = []

        # Configure the election builder
        description = self.get_hamilton_election_from_file()
        builder = ElectionBuilder(NUMBER_OF_GUARDIANS, QUORUM, description)

        # Setup Guardians
        for i in range(NUMBER_OF_GUARDIANS):
            guardians.append(
                Guardian(
                    "hamilton-county-canvass-board-member-" + str(i),
                    i,
                    NUMBER_OF_GUARDIANS,
                    QUORUM,
                ))

        # Run the key ceremony
        mediator = KeyCeremonyMediator(guardians[0].ceremony_details)
        for guardian in guardians:
            mediator.announce(guardian)
        mediator.orchestrate()
        mediator.verify()

        # Joint Key
        joint_key = mediator.publish_joint_key()

        # Save Validation Keys
        for guardian in guardians:
            coefficient_validation_sets.append(
                guardian.share_coefficient_validation_set())

        builder.set_public_key(get_optional(joint_key).joint_public_key)
        builder.set_commitment_hash(get_optional(joint_key).commitment_hash)
        metadata, context = get_optional(builder.build())
        constants = ElectionConstants()

        return (
            AllPublicElectionData(description, metadata, context, constants,
                                  coefficient_validation_sets),
            AllPrivateElectionData(guardians),
        )
class TestDecryptionMediator(TestCase):
    NUMBER_OF_GUARDIANS = 3
    QUORUM = 2
    CEREMONY_DETAILS = CeremonyDetails(NUMBER_OF_GUARDIANS, QUORUM)

    def setUp(self):

        self.key_ceremony = KeyCeremonyMediator(self.CEREMONY_DETAILS)

        self.guardians: List[Guardian] = []

        # Setup Guardians
        for i in range(self.NUMBER_OF_GUARDIANS):
            sequence = i + 2
            self.guardians.append(
                Guardian(
                    "guardian_" + str(sequence),
                    sequence,
                    self.NUMBER_OF_GUARDIANS,
                    self.QUORUM,
                ))

        # Attendance (Public Key Share)
        for guardian in self.guardians:
            self.key_ceremony.announce(guardian)

        self.key_ceremony.orchestrate(identity_auxiliary_encrypt)
        self.key_ceremony.verify(identity_auxiliary_decrypt)

        self.joint_public_key = self.key_ceremony.publish_joint_key()
        self.assertIsNotNone(self.joint_public_key)

        # setup the election
        self.election = election_factory.get_fake_election()
        builder = ElectionBuilder(self.NUMBER_OF_GUARDIANS, self.QUORUM,
                                  self.election)

        self.assertIsNone(
            builder.build())  # Can't build without the public key

        builder.set_public_key(self.joint_public_key)
        self.metadata, self.context = get_optional(builder.build())

        self.encryption_device = EncryptionDevice("location")
        self.ballot_marking_device = EncryptionMediator(
            self.metadata, self.context, self.encryption_device)

        # get some fake ballots
        self.fake_cast_ballot = ballot_factory.get_fake_ballot(
            self.metadata, "some-unique-ballot-id-cast")
        self.more_fake_ballots = []
        for i in range(10):
            self.more_fake_ballots.append(
                ballot_factory.get_fake_ballot(
                    self.metadata, f"some-unique-ballot-id-cast{i}"))
        self.fake_spoiled_ballot = ballot_factory.get_fake_ballot(
            self.metadata, "some-unique-ballot-id-spoiled")
        self.assertTrue(
            self.fake_cast_ballot.is_valid(
                self.metadata.ballot_styles[0].object_id))
        self.assertTrue(
            self.fake_spoiled_ballot.is_valid(
                self.metadata.ballot_styles[0].object_id))
        self.expected_plaintext_tally = accumulate_plaintext_ballots(
            [self.fake_cast_ballot] + self.more_fake_ballots)

        # Fill in the expected values with any missing selections
        # that were not made on any ballots
        selection_ids = set([
            selection.object_id for contest in self.metadata.contests
            for selection in contest.ballot_selections
        ])

        missing_selection_ids = selection_ids.difference(
            set(self.expected_plaintext_tally))

        for id in missing_selection_ids:
            self.expected_plaintext_tally[id] = 0

        # Encrypt
        encrypted_fake_cast_ballot = self.ballot_marking_device.encrypt(
            self.fake_cast_ballot)
        encrypted_fake_spoiled_ballot = self.ballot_marking_device.encrypt(
            self.fake_spoiled_ballot)
        self.assertIsNotNone(encrypted_fake_cast_ballot)
        self.assertIsNotNone(encrypted_fake_spoiled_ballot)
        self.assertTrue(
            encrypted_fake_cast_ballot.is_valid_encryption(
                self.context.crypto_extended_base_hash, self.joint_public_key))

        # encrypt some more fake ballots
        self.more_fake_encrypted_ballots = []
        for fake_ballot in self.more_fake_ballots:
            self.more_fake_encrypted_ballots.append(
                self.ballot_marking_device.encrypt(fake_ballot))

        # configure the ballot box
        ballot_store = BallotStore()
        ballot_box = BallotBox(self.metadata, self.context, ballot_store)
        ballot_box.cast(encrypted_fake_cast_ballot)
        ballot_box.spoil(encrypted_fake_spoiled_ballot)

        # Cast some more fake ballots
        for fake_ballot in self.more_fake_encrypted_ballots:
            ballot_box.cast(fake_ballot)

        # generate encrypted tally
        self.ciphertext_tally = tally_ballots(ballot_store, self.metadata,
                                              self.context)

    def tearDown(self):
        self.key_ceremony.reset(
            CeremonyDetails(self.NUMBER_OF_GUARDIANS, self.QUORUM))

    def test_announce(self):
        # Arrange
        subject = DecryptionMediator(self.metadata, self.context,
                                     self.ciphertext_tally)

        # act
        result = subject.announce(self.guardians[0])

        # assert
        self.assertIsNotNone(result)

        # Can only announce once
        self.assertIsNotNone(subject.announce(self.guardians[0]))

        # Cannot submit another share internally
        self.assertFalse(
            subject._submit_decryption_share(
                TallyDecryptionShare(self.guardians[0].object_id, ZERO_MOD_P,
                                     {}, {})))

        # Cannot get plaintext tally without a quorum
        self.assertIsNone(subject.get_plaintext_tally())

    def test_compute_selection(self):
        # Arrange
        first_selection = [
            selection for contest in self.ciphertext_tally.cast.values()
            for selection in contest.tally_selections.values()
        ][0]

        # act
        result = compute_decryption_share_for_selection(
            self.guardians[0], first_selection, self.context)

        # assert
        self.assertIsNotNone(result)

    def test_compute_compensated_selection_failure(self):
        # Arrange
        first_selection = [
            selection for contest in self.ciphertext_tally.cast.values()
            for selection in contest.tally_selections.values()
        ][0]

        # Act
        self.guardians[0]._guardian_election_partial_key_backups.pop(
            self.guardians[2].object_id)

        self.assertIsNone(self.guardians[0].recovery_public_key_for(
            self.guardians[2].object_id))

        result = compute_compensated_decryption_share_for_selection(
            self.guardians[0],
            self.guardians[2].object_id,
            first_selection,
            self.context,
            identity_auxiliary_decrypt,
        )

        # Assert
        self.assertIsNone(result)

    def test_compute_compensated_selection(self):
        """
        demonstrates the complete workflow for computing a comepnsated decryption share
        For one selection. It is useful for verifying that the workflow is correct
        """
        # Arrange
        first_selection = [
            selection for contest in self.ciphertext_tally.cast.values()
            for selection in contest.tally_selections.values()
        ][0]

        # Compute lagrange coefficients for the guardians that are present
        lagrange_0 = compute_lagrange_coefficient(
            self.guardians[0].sequence_order,
            *[self.guardians[1].sequence_order],
        )
        lagrange_1 = compute_lagrange_coefficient(
            self.guardians[1].sequence_order,
            *[self.guardians[0].sequence_order],
        )

        print(
            f"lagrange: sequence_orders: ({self.guardians[0].sequence_order}, {self.guardians[1].sequence_order}, {self.guardians[2].sequence_order})\n"
        )

        print(lagrange_0)
        print(lagrange_1)

        # compute their shares
        share_0 = compute_decryption_share_for_selection(
            self.guardians[0], first_selection, self.context)

        share_1 = compute_decryption_share_for_selection(
            self.guardians[1], first_selection, self.context)

        self.assertIsNotNone(share_0)
        self.assertIsNotNone(share_1)

        # compute compensations shares for the missing guardian
        compensation_0 = compute_compensated_decryption_share_for_selection(
            self.guardians[0],
            self.guardians[2].object_id,
            first_selection,
            self.context,
            identity_auxiliary_decrypt,
        )

        compensation_1 = compute_compensated_decryption_share_for_selection(
            self.guardians[1],
            self.guardians[2].object_id,
            first_selection,
            self.context,
            identity_auxiliary_decrypt,
        )

        self.assertIsNotNone(compensation_0)
        self.assertIsNotNone(compensation_1)

        print("\nSHARES:")
        print(compensation_0)
        print(compensation_1)

        # Check the share proofs
        self.assertTrue(
            compensation_0.proof.is_valid(
                first_selection.message,
                get_optional(self.guardians[0].recovery_public_key_for(
                    self.guardians[2].object_id)),
                compensation_0.share,
                self.context.crypto_extended_base_hash,
            ))

        self.assertTrue(
            compensation_1.proof.is_valid(
                first_selection.message,
                get_optional(self.guardians[1].recovery_public_key_for(
                    self.guardians[2].object_id)),
                compensation_1.share,
                self.context.crypto_extended_base_hash,
            ))

        share_pow_p = [
            pow_p(compensation_0.share, lagrange_0),
            pow_p(compensation_1.share, lagrange_1),
        ]

        print("\nSHARE_POW_P")
        print(share_pow_p)

        # reconstruct the missing share from the compensation shares
        reconstructed_share = mult_p(*[
            pow_p(compensation_0.share, lagrange_0),
            pow_p(compensation_1.share, lagrange_1),
        ])

        print("\nRECONSTRUCTED SHARE\n")
        print(reconstructed_share)

        share_2 = CiphertextDecryptionSelection(
            first_selection.object_id,
            self.guardians[2].object_id,
            first_selection.description_hash,
            reconstructed_share,
            {
                self.guardians[0].object_id: compensation_0,
                self.guardians[1].object_id: compensation_1,
            },
        )

        # Decrypt the result
        result = decrypt_selection_with_decryption_shares(
            first_selection,
            {
                self.guardians[0].object_id: (
                    self.guardians[0].share_election_public_key().key,
                    share_0,
                ),
                self.guardians[1].object_id: (
                    self.guardians[1].share_election_public_key().key,
                    share_1,
                ),
                self.guardians[2].object_id: (
                    self.guardians[2].share_election_public_key().key,
                    share_2,
                ),
            },
            self.context.crypto_extended_base_hash,
        )

        print(result)

        self.assertIsNotNone(result)
        self.assertEqual(
            result.plaintext,
            self.expected_plaintext_tally[first_selection.object_id])

    def test_decrypt_selection_all_present(self):
        # Arrange

        # find the first selection
        first_contest = [
            contest for contest in self.ciphertext_tally.cast.values()
        ][0]
        first_selection = list(first_contest.tally_selections.values())[0]

        # precompute decryption shares for the guardians
        first_share = compute_decryption_share(self.guardians[0],
                                               self.ciphertext_tally,
                                               self.context)
        second_share = compute_decryption_share(self.guardians[1],
                                                self.ciphertext_tally,
                                                self.context)
        third_share = compute_decryption_share(self.guardians[2],
                                               self.ciphertext_tally,
                                               self.context)

        # build type: Dict[GUARDIAN_ID, Tuple[ELECTION_PUBLIC_KEY, TallyDecryptionShare]]
        shares = {
            self.guardians[0].object_id: (
                self.guardians[0].share_election_public_key().key,
                first_share.contests[first_contest.object_id].selections[
                    first_selection.object_id],
            ),
            self.guardians[1].object_id: (
                self.guardians[1].share_election_public_key().key,
                second_share.contests[first_contest.object_id].selections[
                    first_selection.object_id],
            ),
            self.guardians[2].object_id: (
                self.guardians[2].share_election_public_key().key,
                third_share.contests[first_contest.object_id].selections[
                    first_selection.object_id],
            ),
        }

        # act
        result = decrypt_selection_with_decryption_shares(
            first_selection, shares, self.context.crypto_extended_base_hash)

        # assert
        self.assertIsNotNone(result)
        self.assertEqual(
            self.expected_plaintext_tally[first_selection.object_id],
            result.plaintext)

    def test_decrypt_spoiled_ballots_all_guardians_present(self):
        # Arrange
        # precompute decryption shares for the guardians
        first_share = compute_decryption_share(self.guardians[0],
                                               self.ciphertext_tally,
                                               self.context)
        second_share = compute_decryption_share(self.guardians[1],
                                                self.ciphertext_tally,
                                                self.context)
        third_share = compute_decryption_share(self.guardians[2],
                                               self.ciphertext_tally,
                                               self.context)
        shares = {
            self.guardians[0].object_id: first_share,
            self.guardians[1].object_id: second_share,
            self.guardians[2].object_id: third_share,
        }

        subject = DecryptionMediator(self.metadata, self.context,
                                     self.ciphertext_tally)

        # act
        result = decrypt_spoiled_ballots(
            self.ciphertext_tally.spoiled_ballots,
            shares,
            self.context.crypto_extended_base_hash,
        )

        # assert
        self.assertIsNotNone(result)
        self.assertTrue(self.fake_spoiled_ballot.object_id in result)

        spoiled_ballot = result[self.fake_spoiled_ballot.object_id]
        for contest in self.fake_spoiled_ballot.contests:
            for selection in contest.ballot_selections:
                self.assertEqual(
                    spoiled_ballot[contest.object_id].selections[
                        selection.object_id].plaintext,
                    result[self.fake_spoiled_ballot.object_id][
                        contest.object_id].selections[
                            selection.object_id].plaintext,
                )

    def test_get_plaintext_tally_all_guardians_present_simple(self):
        # Arrange
        subject = DecryptionMediator(self.metadata, self.context,
                                     self.ciphertext_tally)

        # act
        for guardian in self.guardians:
            self.assertIsNotNone(subject.announce(guardian))

        decrypted_tallies = subject.get_plaintext_tally()
        result = self._convert_to_selections(decrypted_tallies)

        # assert
        self.assertIsNotNone(result)
        self.assertEqual(self.expected_plaintext_tally, result)

        # Verify we get the same tally back if we call again
        another_decrypted_tally = subject.get_plaintext_tally()

        self.assertEqual(decrypted_tallies, another_decrypted_tally)

    def test_get_plaintext_tally_compensate_missing_guardian_simple(self):

        # Arrange
        subject = DecryptionMediator(self.metadata, self.context,
                                     self.ciphertext_tally)

        # Act

        self.assertIsNotNone(subject.announce(self.guardians[0]))
        self.assertIsNotNone(subject.announce(self.guardians[1]))

        # explicitly compensate to demonstrate that this is possible, but not required
        self.assertIsNotNone(
            subject.compensate(self.guardians[2].object_id,
                               identity_auxiliary_decrypt))

        decrypted_tallies = subject.get_plaintext_tally()
        self.assertIsNotNone(decrypted_tallies)
        result = self._convert_to_selections(decrypted_tallies)

        # assert
        self.assertIsNotNone(result)
        print(result)
        self.assertEqual(self.expected_plaintext_tally, result)

    @settings(
        deadline=timedelta(milliseconds=15000),
        suppress_health_check=[HealthCheck.too_slow],
        max_examples=8,
        # disabling the "shrink" phase, because it runs very slowly
        phases=[Phase.explicit, Phase.reuse, Phase.generate, Phase.target],
    )
    @given(data(), integers(1, 3), integers(2, 5))
    def test_get_plaintext_tally_all_guardians_present(self, data,
                                                       parties: int,
                                                       contests: int):
        # Arrange
        description = data.draw(election_descriptions(parties, contests))
        builder = ElectionBuilder(self.NUMBER_OF_GUARDIANS, self.QUORUM,
                                  description)
        metadata, context = builder.set_public_key(
            self.joint_public_key).build()

        plaintext_ballots: List[PlaintextBallot] = data.draw(
            plaintext_voted_ballots(metadata, randrange(3, 6)))
        plaintext_tallies = accumulate_plaintext_ballots(plaintext_ballots)

        encrypted_tally = self._generate_encrypted_tally(
            metadata, context, plaintext_ballots)

        subject = DecryptionMediator(metadata, context, encrypted_tally)

        # act
        for guardian in self.guardians:
            self.assertIsNotNone(subject.announce(guardian))

        decrypted_tallies = subject.get_plaintext_tally()
        result = self._convert_to_selections(decrypted_tallies)

        # assert
        self.assertIsNotNone(result)
        self.assertEqual(plaintext_tallies, result)

    def _generate_encrypted_tally(
        self,
        metadata: InternalElectionDescription,
        context: CiphertextElectionContext,
        ballots: List[PlaintextBallot],
    ) -> CiphertextTally:

        # encrypt each ballot
        store = BallotStore()
        for ballot in ballots:
            encrypted_ballot = encrypt_ballot(ballot, metadata, context,
                                              int_to_q_unchecked(1))
            self.assertIsNotNone(encrypted_ballot)
            # add to the ballot store
            store.set(
                encrypted_ballot.object_id,
                from_ciphertext_ballot(encrypted_ballot, BallotBoxState.CAST),
            )

        tally = tally_ballots(store, metadata, context)
        self.assertIsNotNone(tally)
        return get_optional(tally)

    def _convert_to_selections(self, tally: PlaintextTally) -> Dict[str, int]:
        plaintext_selections: Dict[str, int] = {}
        for _, contest in tally.contests.items():
            for selection_id, selection in contest.selections.items():
                plaintext_selections[selection_id] = selection.plaintext

        return plaintext_selections
Exemple #4
0
class TestEndToEndElection(TestCase):
    """
    Test a complete simple example of executing an End-to-End encrypted election.
    In a real world scenario all of these steps would not be completed on the same machine.
    """

    NUMBER_OF_GUARDIANS = 5
    QUORUM = 3

    REMOVE_OUTPUT = False

    # Step 0 - Configure Election
    description: ElectionDescription
    election_builder: ElectionBuilder
    metadata: InternalElectionDescription
    context: CiphertextElectionContext
    constants: ElectionConstants

    # Step 1 - Key Ceremony
    mediator: KeyCeremonyMediator
    guardians: List[Guardian] = []
    coefficient_validation_sets: List[CoefficientValidationSet] = []

    # Step 2 - Encrypt Votes
    device: EncryptionDevice
    encrypter: EncryptionMediator
    plaintext_ballots: List[PlaintextBallot]
    ciphertext_ballots: List[CiphertextBallot] = []

    # Step 3 - Cast and Spoil
    ballot_store: BallotStore
    ballot_box: BallotBox

    # Step 4 - Decrypt Tally
    ciphertext_tally: CiphertextTally
    plaintext_tally: PlaintextTally
    decrypter: DecryptionMediator

    def test_end_to_end_election(self) -> None:
        """
        Execute the simplified end-to-end test demonstrating each component of the system.
        """
        self.step_0_configure_election()
        self.step_1_key_ceremony()
        self.step_2_encrypt_votes()
        self.step_3_cast_and_spoil()
        self.step_4_decrypt_tally()
        self.step_5_publish_and_verify()

    def step_0_configure_election(self) -> None:
        """
        To conduct an election, load an `ElectionDescription` file
        """

        # Load a pre-configured Election Description
        # TODO: replace with complex election
        self.description = ElectionFactory().get_simple_election_from_file()
        print(f"""
            {'-'*40}\n
            # Election Summary:
            # Scope: {self.description.election_scope_id}
            # Geopolitical Units: {len(self.description.geopolitical_units)}
            # Parties: {len(self.description.parties)}
            # Candidates: {len(self.description.candidates)}
            # Contests: {len(self.description.contests)}
            # Ballot Styles: {len(self.description.ballot_styles)}\n
            {'-'*40}\n
            """)
        self._assert_message(
            ElectionDescription.is_valid.__qualname__,
            "Verify that the input election meta-data is well-formed",
            self.description.is_valid(),
        )

        # Create an Election Builder
        self.election_builder = ElectionBuilder(self.NUMBER_OF_GUARDIANS,
                                                self.QUORUM, self.description)
        self._assert_message(
            ElectionBuilder.__qualname__,
            f"Created with number_of_guardians: {self.NUMBER_OF_GUARDIANS} quorum: {self.QUORUM}",
        )

        # Move on to the Key Ceremony

    def step_1_key_ceremony(self) -> None:
        """
        Using the NUMBER_OF_GUARDIANS, generate public-private keypairs and share
        representations of those keys with QUORUM of other Guardians.  Then, combine
        the public election keys to make a joint election key that is used to encrypt ballots
        """

        # Setup Guardians
        for i in range(self.NUMBER_OF_GUARDIANS):
            self.guardians.append(
                Guardian("guardian_" + str(i), i, self.NUMBER_OF_GUARDIANS,
                         self.QUORUM))

        # Setup Mediator
        self.mediator = KeyCeremonyMediator(self.guardians[0].ceremony_details)

        # Attendance (Public Key Share)
        for guardian in self.guardians:
            self.mediator.announce(guardian)

        self._assert_message(
            KeyCeremonyMediator.all_guardians_in_attendance.__qualname__,
            "Confirms all guardians have shared their public keys",
            self.mediator.all_guardians_in_attendance(),
        )

        # Run the Key Ceremony process,
        # Which shares the keys among the guardians
        orchestrated = self.mediator.orchestrate()
        self._assert_message(
            KeyCeremonyMediator.orchestrate.__qualname__,
            "Executes the key exchange between guardians",
            orchestrated is not None,
        )

        self._assert_message(
            KeyCeremonyMediator.all_election_partial_key_backups_available.
            __qualname__,
            "Confirm sall guardians have shared their partial key backups",
            self.mediator.all_election_partial_key_backups_available(),
        )

        # Verification
        verified = self.mediator.verify()
        self._assert_message(
            KeyCeremonyMediator.verify.__qualname__,
            "Confirms all guardians truthfully executed the ceremony",
            verified,
        )

        self._assert_message(
            KeyCeremonyMediator.
            all_election_partial_key_verifications_received.__qualname__,
            "Confirms all guardians have submitted a verification of the backups of all other guardians",
            self.mediator.all_election_partial_key_verifications_received(),
        )

        self._assert_message(
            KeyCeremonyMediator.all_election_partial_key_backups_verified.
            __qualname__,
            "Confirms all guardians have verified the backups of all other guardians",
            self.mediator.all_election_partial_key_backups_verified(),
        )

        # Joint Key
        joint_key = self.mediator.publish_joint_key()
        self._assert_message(
            KeyCeremonyMediator.publish_joint_key.__qualname__,
            "Publishes the Joint Election Key",
            joint_key is not None,
        )

        # Save Validation Keys
        for guardian in self.guardians:
            self.coefficient_validation_sets.append(
                guardian.share_coefficient_validation_set())

        # Build the Election
        self.election_builder.set_public_key(get_optional(joint_key))
        self.metadata, self.context = get_optional(
            self.election_builder.build())
        self.constants = ElectionConstants()

        # Move on to encrypting ballots

    def step_2_encrypt_votes(self) -> None:
        """
        Using the `CiphertextElectionContext` encrypt ballots for the election
        """

        # Configure the Encryption Device
        self.device = EncryptionDevice("polling-place-one")
        self.encrypter = EncryptionMediator(self.metadata, self.context,
                                            self.device)
        self._assert_message(
            EncryptionDevice.__qualname__,
            f"Ready to encrypt at location: {self.device.location}",
        )

        # Load some Ballots
        self.plaintext_ballots = BallotFactory().get_simple_ballots_from_file()
        self._assert_message(
            PlaintextBallot.__qualname__,
            f"Loaded ballots: {len(self.plaintext_ballots)}",
            len(self.plaintext_ballots) > 0,
        )

        # Encrypt the Ballots
        for plaintext_ballot in self.plaintext_ballots:
            encrypted_ballot = self.encrypter.encrypt(plaintext_ballot)
            self._assert_message(
                EncryptionMediator.encrypt.__qualname__,
                f"Ballot Id: {plaintext_ballot.object_id}",
                encrypted_ballot is not None,
            )
            self.ciphertext_ballots.append(get_optional(encrypted_ballot))

        # Next, we cast or spoil the ballots

    def step_3_cast_and_spoil(self) -> None:
        """
        Accept each ballot by marking it as either cast or spoiled.
        This example demonstrates one way to accept ballots using the `BallotBox` class
        """

        # Configure the Ballot Box
        self.ballot_store = BallotStore()
        self.ballot_box = BallotBox(self.metadata, self.context,
                                    self.ballot_store)

        # Randomly cast or spoil the ballots
        for ballot in self.ciphertext_ballots:
            # if randint(0, 1):
            #     accepted_ballot = self.ballot_box.cast(ballot)
            # else:
            #     accepted_ballot = self.ballot_box.spoil(ballot)
            accepted_ballot = self.ballot_box.cast(ballot)

            self._assert_message(
                BallotBox.__qualname__,
                f"Accepted Ballot Id: {ballot.object_id} state: {get_optional(accepted_ballot).state}",
                accepted_ballot is not None,
            )

    def step_4_decrypt_tally(self) -> None:
        """
        Homomorphically combine the selections made on all of the cast ballots
        and use the Available Guardians to decrypt the combined tally.
        In this way, no individual voter's cast ballot is ever decrypted drectly.
        """

        # Generate a Homomorphically Accumulated Tally of the ballots
        self.ciphertext_tally = get_optional(
            tally_ballots(self.ballot_store, self.metadata, self.context))
        self._assert_message(
            tally_ballots.__qualname__,
            f"""
            - cast: {self.ciphertext_tally.count()} 
            - spoiled: {len(self.ciphertext_tally.spoiled_ballots)}
            Total: {len(self.ciphertext_tally)}
            """,
            self.ciphertext_tally is not None,
        )

        # Configure the Decryption
        self.decrypter = DecryptionMediator(self.metadata, self.context,
                                            self.ciphertext_tally)

        # Announce each guardian as present
        for guardian in self.guardians:
            decryption_share = self.decrypter.announce(guardian)
            self._assert_message(
                DecryptionMediator.announce.__qualname__,
                f"Guardian Present: {guardian.object_id}",
                decryption_share is not None,
            )

        # Get the Plain Text Tally
        self.plaintext_tally = get_optional(
            self.decrypter.get_plaintext_tally())
        self._assert_message(
            DecryptionMediator.get_plaintext_tally.__qualname__,
            "Tally Decrypted",
            self.plaintext_tally is not None,
        )

        # Now, compare the results
        self.compare_results()

    def compare_results(self) -> None:
        """
        Compare the results to ensure the decryption was done correctly
        """
        print(f"""
            {'-'*40}
            # Election Results:

            """)

        # Create a representation of each contest's tally
        selection_ids = [
            selection.object_id for contest in self.metadata.contests
            for selection in contest.ballot_selections
        ]
        expected_plaintext_tally: Dict[str, int] = {
            key: 0
            for key in selection_ids
        }

        # Tally the expected values from the loaded ballots
        for ballot in self.plaintext_ballots:
            if (get_optional(self.ballot_store.get(
                    ballot.object_id)).state == BallotBoxState.CAST):
                for contest in ballot.contests:
                    for selection in contest.ballot_selections:
                        expected_plaintext_tally[
                            selection.object_id] += selection.to_int()

        # Compare the expected tally to the decrypted tally
        for tally_contest in self.plaintext_tally.contests.values():
            print(f" Contest: {tally_contest.object_id}")
            for tally_selection in tally_contest.selections.values():
                expected = expected_plaintext_tally[tally_selection.object_id]
                self._assert_message(
                    f"   - Selection: {tally_selection.object_id}",
                    f"expected: {expected}, actual: {tally_selection.tally}",
                    expected == tally_selection.tally,
                )
        print(f"\n{'-'*40}\n")

        # Compare the expected values for each spoiled ballot
        for ballot_id, accepted_ballot in self.ciphertext_tally.spoiled_ballots.items(
        ):
            if accepted_ballot.state == BallotBoxState.SPOILED:
                for plaintext_ballot in self.plaintext_ballots:
                    if ballot_id == plaintext_ballot.object_id:
                        print(f"\nSpoiled Ballot: {ballot_id}")
                        for contest in plaintext_ballot.contests:
                            print(f"\n Contest: {contest.object_id}")
                            for selection in contest.ballot_selections:
                                expected = selection.to_int()
                                decrypted_selection = (
                                    self.plaintext_tally.spoiled_ballots[
                                        ballot_id][contest.object_id].
                                    selections[selection.object_id])
                                self._assert_message(
                                    f"   - Selection: {selection.object_id}",
                                    f"expected: {expected}, actual: {decrypted_selection.tally}",
                                    expected == decrypted_selection.tally,
                                )

    def step_5_publish_and_verify(self) -> None:
        """Publish and verify steps of the election"""
        self.publish_results()
        self.verify_results()

        if self.REMOVE_OUTPUT:
            rmtree(RESULTS_DIR)

    def publish_results(self) -> None:
        """
        Publish results/artifacts of the election
        """
        publish(
            self.description,
            self.context,
            self.constants,
            [self.device],
            self.ballot_store.all(),
            self.ciphertext_tally.spoiled_ballots.values(),
            publish_ciphertext_tally(self.ciphertext_tally),
            self.plaintext_tally,
            self.coefficient_validation_sets,
        )
        self._assert_message(
            "Publish",
            f"Artifacts published to: {RESULTS_DIR}",
            path.exists(RESULTS_DIR),
        )

    def verify_results(self) -> None:
        """Verify results of election"""

        # Deserialize
        description_from_file = ElectionDescription.from_json_file(
            DESCRIPTION_FILE_NAME, RESULTS_DIR)
        self.assertEqual(self.description, description_from_file)

        context_from_file = CiphertextElectionContext.from_json_file(
            CONTEXT_FILE_NAME, RESULTS_DIR)
        self.assertEqual(self.context, context_from_file)

        constants_from_file = ElectionConstants.from_json_file(
            CONSTANTS_FILE_NAME, RESULTS_DIR)
        self.assertEqual(self.constants, constants_from_file)

        device_name = DEVICE_PREFIX + str(self.device.uuid)
        device_from_file = EncryptionDevice.from_json_file(
            device_name, DEVICES_DIR)
        self.assertEqual(self.device, device_from_file)

        ciphertext_ballots: List[CiphertextAcceptedBallot] = []
        for ballot in self.ballot_store.all():
            ballot_name = BALLOT_PREFIX + ballot.object_id
            ballot_from_file = CiphertextAcceptedBallot.from_json_file(
                ballot_name, BALLOTS_DIR)
            self.assertEqual(ballot, ballot_from_file)

        spoiled_ballots: List[CiphertextAcceptedBallot] = []
        for spoiled_ballot in self.ciphertext_tally.spoiled_ballots.values():
            ballot_name = BALLOT_PREFIX + spoiled_ballot.object_id
            spoiled_ballot_from_file = CiphertextAcceptedBallot.from_json_file(
                ballot_name, SPOILED_DIR)
            self.assertEqual(spoiled_ballot, spoiled_ballot_from_file)

        ciphertext_tally_from_file = PublishedCiphertextTally.from_json_file(
            ENCRYPTED_TALLY_FILE_NAME, RESULTS_DIR)
        self.assertEqual(publish_ciphertext_tally(self.ciphertext_tally),
                         ciphertext_tally_from_file)

        plainttext_tally_from_file = PlaintextTally.from_json_file(
            TALLY_FILE_NAME, RESULTS_DIR)
        self.assertEqual(self.plaintext_tally, plainttext_tally_from_file)

        coefficient_validation_sets: List[CoefficientValidationSet] = []
        for coefficient_validation_set in self.coefficient_validation_sets:
            set_name = COEFFICIENT_PREFIX + coefficient_validation_set.owner_id
            coefficient_validation_set_from_file = (
                CoefficientValidationSet.from_json_file(
                    set_name, COEFFICIENTS_DIR))
            self.assertEqual(coefficient_validation_set,
                             coefficient_validation_set_from_file)

    def _assert_message(self,
                        name: str,
                        message: str,
                        condition: Union[Callable, bool] = True) -> None:
        if callable(condition):
            result = condition()
        else:
            result = condition

        print(f"{name}: {message}: {result}")
        self.assertTrue(result)
class TestDecryptionMediator(TestCase):
    NUMBER_OF_GUARDIANS = 3
    QUORUM = 2
    CEREMONY_DETAILS = CeremonyDetails(NUMBER_OF_GUARDIANS, QUORUM)

    def setUp(self):

        self.key_ceremony = KeyCeremonyMediator(self.CEREMONY_DETAILS)

        self.guardians: List[Guardian] = []

        # Setup Guardians
        for i in range(self.NUMBER_OF_GUARDIANS):
            self.guardians.append(
                Guardian("guardian_" + str(i), i, self.NUMBER_OF_GUARDIANS, self.QUORUM)
            )

        # Attendance (Public Key Share)
        for guardian in self.guardians:
            self.key_ceremony.announce(guardian)

        self.key_ceremony.orchestrate()
        self.key_ceremony.verify()

        self.joint_public_key = self.key_ceremony.publish_joint_key()
        self.assertIsNotNone(self.joint_public_key)

        # setup the election
        self.election = election_factory.get_fake_election()
        builder = ElectionBuilder(self.NUMBER_OF_GUARDIANS, self.QUORUM, self.election)

        self.assertIsNone(builder.build())  # Can't build without the public key

        builder.set_public_key(self.joint_public_key)
        self.metadata, self.context = get_optional(builder.build())

        self.encryption_device = EncryptionDevice("location")
        self.ballot_marking_device = EncryptionMediator(
            self.metadata, self.context, self.encryption_device
        )

        # get some fake ballots
        self.fake_cast_ballot = ballot_factory.get_fake_ballot(
            self.metadata, "some-unique-ballot-id-cast"
        )
        self.fake_spoiled_ballot = ballot_factory.get_fake_ballot(
            self.metadata, "some-unique-ballot-id-spoiled"
        )
        self.assertTrue(
            self.fake_cast_ballot.is_valid(self.metadata.ballot_styles[0].object_id)
        )
        self.assertTrue(
            self.fake_spoiled_ballot.is_valid(self.metadata.ballot_styles[0].object_id)
        )
        self.expected_plaintext_tally = accumulate_plaintext_ballots(
            [self.fake_cast_ballot]
        )

        # Fill in any missing selections that were not made on any ballots
        selection_ids = set(
            [
                selection.object_id
                for contest in self.metadata.contests
                for selection in contest.ballot_selections
            ]
        )

        missing_selection_ids = selection_ids.difference(
            set(self.expected_plaintext_tally)
        )

        for id in missing_selection_ids:
            self.expected_plaintext_tally[id] = 0

        # Encrypt
        encrypted_fake_cast_ballot = self.ballot_marking_device.encrypt(
            self.fake_cast_ballot
        )
        encrypted_fake_spoiled_ballot = self.ballot_marking_device.encrypt(
            self.fake_spoiled_ballot
        )
        self.assertIsNotNone(encrypted_fake_cast_ballot)
        self.assertIsNotNone(encrypted_fake_spoiled_ballot)
        self.assertTrue(
            encrypted_fake_cast_ballot.is_valid_encryption(
                self.context.crypto_extended_base_hash, self.joint_public_key
            )
        )

        # configure the ballot box
        ballot_store = BallotStore()
        ballot_box = BallotBox(self.metadata, self.context, ballot_store)
        ballot_box.cast(encrypted_fake_cast_ballot)
        ballot_box.spoil(encrypted_fake_spoiled_ballot)

        # generate encrypted tally
        self.ciphertext_tally = tally_ballots(ballot_store, self.metadata, self.context)

    def test_announce(self):
        # Arrange
        subject = DecryptionMediator(self.metadata, self.context, self.ciphertext_tally)

        # act
        result = subject.announce(self.guardians[0])

        # assert
        self.assertIsNotNone(result)

        # Can only announce once
        self.assertIsNone(subject.announce(self.guardians[0]))

        # Cannot submit another share internally
        self.assertFalse(
            subject._submit_decryption_share(
                DecryptionShare(self.guardians[0].object_id, {}, {})
            )
        )

        # Cannot get plaintext tally without a quorum
        self.assertIsNone(subject.get_plaintext_tally())

    def test_compute_selection(self):
        # Arrange
        first_selection = [
            selection
            for contest in self.ciphertext_tally.cast.values()
            for selection in contest.tally_selections.values()
        ][0]

        # act
        result = _compute_decryption_for_selection(
            self.guardians[0], first_selection, self.context
        )

        # assert
        self.assertIsNotNone(result)

    def test_decrypt_selection(self):
        # Arrange

        # find the first selection
        first_contest = [contest for contest in self.ciphertext_tally.cast.values()][0]
        first_selection = list(first_contest.tally_selections.values())[0]

        # precompute decryption shares for the guardians
        first_share = compute_decryption_share(
            self.guardians[0], self.ciphertext_tally, self.context
        )
        second_share = compute_decryption_share(
            self.guardians[1], self.ciphertext_tally, self.context
        )
        third_share = compute_decryption_share(
            self.guardians[2], self.ciphertext_tally, self.context
        )
        shares = {
            self.guardians[0]
            .object_id: first_share.contests[first_contest.object_id]
            .selections[first_selection.object_id]
            .share,
            self.guardians[1]
            .object_id: second_share.contests[first_contest.object_id]
            .selections[first_selection.object_id]
            .share,
            self.guardians[2]
            .object_id: third_share.contests[first_contest.object_id]
            .selections[first_selection.object_id]
            .share,
        }

        # act
        result = decrypt_selection_with_decryption_shares(first_selection, shares)

        # assert
        self.assertIsNotNone(result)
        self.assertEqual(
            self.expected_plaintext_tally[first_selection.object_id], result.plaintext
        )

    def test_decrypt_spoiled_ballots(self):
        # Arrange
        # precompute decryption shares for the guardians
        first_share = compute_decryption_share(
            self.guardians[0], self.ciphertext_tally, self.context
        )
        second_share = compute_decryption_share(
            self.guardians[1], self.ciphertext_tally, self.context
        )
        third_share = compute_decryption_share(
            self.guardians[2], self.ciphertext_tally, self.context
        )
        shares = {
            self.guardians[0].object_id: first_share,
            self.guardians[1].object_id: second_share,
            self.guardians[2].object_id: third_share,
        }

        subject = DecryptionMediator(self.metadata, self.context, self.ciphertext_tally)

        # act
        result = subject._decrypt_spoiled_ballots(
            self.ciphertext_tally.spoiled_ballots, shares
        )

        # assert
        self.assertIsNotNone(result)
        self.assertTrue(self.fake_spoiled_ballot.object_id in result)

        spoiled_ballot = result[self.fake_spoiled_ballot.object_id]
        for contest in self.fake_spoiled_ballot.contests:
            for selection in contest.ballot_selections:
                self.assertEqual(
                    spoiled_ballot[contest.object_id]
                    .selections[selection.object_id]
                    .plaintext,
                    result[self.fake_spoiled_ballot.object_id][contest.object_id]
                    .selections[selection.object_id]
                    .plaintext,
                )

    def test_get_plaintext_tally_all_guardians_present_simple(self):
        # Arrange
        subject = DecryptionMediator(self.metadata, self.context, self.ciphertext_tally)

        # act
        for guardian in self.guardians:
            self.assertIsNotNone(subject.announce(guardian))

        decrypted_tallies = subject.get_plaintext_tally()
        result = self._convert_to_selections(decrypted_tallies)

        # assert
        self.assertIsNotNone(result)
        self.assertEqual(self.expected_plaintext_tally, result)

    @settings(
        deadline=timedelta(milliseconds=10000),
        suppress_health_check=[HealthCheck.too_slow],
        max_examples=1,
        # disabling the "shrink" phase, because it runs very slowly
        phases=[Phase.explicit, Phase.reuse, Phase.generate, Phase.target],
    )
    @given(data(), integers(1, 3), integers(2, 5))
    def test_get_plaintext_tally_all_guardians_present(
        self, data, parties: int, contests: int
    ):
        # Arrange
        description = data.draw(election_descriptions(parties, contests))
        builder = ElectionBuilder(self.NUMBER_OF_GUARDIANS, self.QUORUM, description)
        metadata, context = builder.set_public_key(self.joint_public_key).build()

        plaintext_ballots: List[PlaintextBallot] = data.draw(
            plaintext_voted_ballots(metadata, randrange(3, 6))
        )
        plaintext_tallies = accumulate_plaintext_ballots(plaintext_ballots)

        encrypted_tally = self._generate_encrypted_tally(
            metadata, context, plaintext_ballots
        )

        subject = DecryptionMediator(metadata, context, encrypted_tally)

        # act
        for guardian in self.guardians:
            self.assertIsNotNone(subject.announce(guardian))

        decrypted_tallies = subject.get_plaintext_tally()
        result = self._convert_to_selections(decrypted_tallies)

        # assert
        self.assertIsNotNone(result)
        self.assertEqual(plaintext_tallies, result)

    def _generate_encrypted_tally(
        self,
        metadata: InternalElectionDescription,
        context: CiphertextElectionContext,
        ballots: List[PlaintextBallot],
    ) -> CiphertextTally:

        # encrypt each ballot
        store = BallotStore()
        for ballot in ballots:
            encrypted_ballot = encrypt_ballot(
                ballot, metadata, context, int_to_q_unchecked(1)
            )
            self.assertIsNotNone(encrypted_ballot)
            # add to the ballot store
            store.set(
                encrypted_ballot.object_id,
                from_ciphertext_ballot(encrypted_ballot, BallotBoxState.CAST),
            )

        tally = tally_ballots(store, metadata, context)
        self.assertIsNotNone(tally)
        return get_optional(tally)

    def _convert_to_selections(self, tally: PlaintextTally) -> Dict[str, int]:
        plaintext_selections: Dict[str, int] = {}
        for _, contest in tally.contests.items():
            for selection_id, selection in contest.selections.items():
                plaintext_selections[selection_id] = selection.plaintext

        return plaintext_selections