✨ Implement Compact Ballot (#331)

* ✨ Implement Compact Ballot

- Add Compact Plaintext and Ciphertext Ballot
- Update Encryption Device and hashing mechanisms
- Isolate `create_ballot_hash` to get a single ballot hash
- Isolate `encrypt_ballot_contests` for using to encrypt all the ballot contests
- Add Compress Methods

* 🚧 Update for Changes in EncryptionDevice

🔨Fix Tracker Test


Squish

* ⚗️ Add Compact Ballot Tests

* ♻️ Refactor ElectionFactory to have Encryption Device

Author: keithrfung

docs/2_Encrypt_Ballots.md

@@ -28,7 +28,7 @@ context: CiphertextElectionContext
 ballot: PlaintextBallot
 
 # Configure an encryption device
-device = EncryptionDevice("polling-place-one")
+device = EncryptionDevice(generate_device_uuid(), "Session", 12345, "polling-place-one")
 encrypter = EncryptionMediator(internal_manifest, context, device)
 
 # Encrypt the ballot

src/electionguard/ballot.py

@@ -871,8 +871,7 @@ def make_ciphertext_ballot(
     if len(contests) == 0:
         log_warning("ciphertext ballot with no contests")
 
-    contest_hashes = [contest.crypto_hash for contest in contests]
-    contest_hash = hash_elems(object_id, manifest_hash, *contest_hashes)
+    contest_hash = create_ballot_hash(object_id, manifest_hash, contests)
 
     timestamp = to_ticks(datetime.now()) if timestamp is None else timestamp
     if previous_ballot_code is None:
@@ -895,6 +894,17 @@ def make_ciphertext_ballot(
     )
 
 
+def create_ballot_hash(
+    ballot_id: str,
+    description_hash: ElementModQ,
+    contests: List[CiphertextBallotContest],
+) -> ElementModQ:
+    """Create the hash of the ballot contests"""
+
+    contest_hashes = [contest.crypto_hash for contest in contests]
+    return hash_elems(ballot_id, description_hash, *contest_hashes)
+
+
 def make_ciphertext_submitted_ballot(
     object_id: str,
     style_id: str,

src/electionguard/ballot_code.py

@@ -2,14 +2,16 @@
 from .group import ElementModQ
 
 
-def get_hash_for_device(uuid: int, location: str) -> ElementModQ:
+def get_hash_for_device(
+    uuid: int, session_id: str, launch_code: int, location: str
+) -> ElementModQ:
     """
     Get starting hash for given device
     :param uuid: Unique identifier of device
     :param location: Location of device
     :return: Starting hash of device
     """
-    return hash_elems(uuid, location)
+    return hash_elems(uuid, session_id, launch_code, location)
 
 
 def get_rotating_ballot_code(

src/electionguard/ballot_compact.py

@@ -0,0 +1,181 @@
+from dataclasses import dataclass
+from typing import Dict, List
+
+
+from .ballot import (
+    SubmittedBallot,
+    ExtendedData,
+    PlaintextBallot,
+    PlaintextBallotContest,
+    PlaintextBallotSelection,
+    create_ballot_hash,
+)
+from .ballot_box import BallotBoxState
+from .election import CiphertextElectionContext
+from .encrypt import encrypt_ballot_contests
+from .group import ElementModQ
+from .manifest import (
+    ContestDescriptionWithPlaceholders,
+    InternalManifest,
+)
+from .utils import get_optional
+
+YES_VOTE = 1
+NO_VOTE = 0
+
+
+@dataclass
+class CompactPlaintextBallot:
+    """A compact plaintext representation of ballot minimized for data size"""
+
+    object_id: str
+    style_id: str
+    selections: List[bool]
+    extended_data: Dict[int, ExtendedData]
+
+
+@dataclass
+class CompactSubmittedBallot:
+    """A compact submitted ballot minimized for data size"""
+
+    compact_plaintext_ballot: CompactPlaintextBallot
+    timestamp: int
+    ballot_nonce: ElementModQ
+    previous_code: ElementModQ
+    code: ElementModQ
+    ballot_box_state: BallotBoxState
+
+
+def compress_plaintext_ballot(ballot: PlaintextBallot) -> CompactPlaintextBallot:
+    """Compress a plaintext ballot into a compact plaintext ballot"""
+    selections = _get_compact_selections(ballot)
+    extended_data = _get_compact_extended_data(ballot)
+
+    return CompactPlaintextBallot(
+        ballot.object_id, ballot.style_id, selections, extended_data
+    )
+
+
+def compress_submitted_ballot(
+    ballot: SubmittedBallot,
+    plaintext_ballot: PlaintextBallot,
+    ballot_nonce: ElementModQ,
+) -> CompactSubmittedBallot:
+    """Compress a submitted ballot into a compact submitted ballot"""
+    return CompactSubmittedBallot(
+        compress_plaintext_ballot(plaintext_ballot),
+        ballot.timestamp,
+        ballot_nonce,
+        ballot.previous_code,
+        get_optional(ballot.code),
+        ballot.state,
+    )
+
+
+def expand_compact_submitted_ballot(
+    compact_ballot: CompactSubmittedBallot,
+    internal_manifest: InternalManifest,
+    context: CiphertextElectionContext,
+) -> SubmittedBallot:
+    """
+    Expand a compact submitted ballot using context and
+    the election manifest into a submitted ballot
+    """
+    # Expand ballot and encrypt & hash contests
+    plaintext_ballot = expand_compact_plaintext_ballot(
+        compact_ballot.compact_plaintext_ballot, internal_manifest
+    )
+    contests = get_optional(
+        encrypt_ballot_contests(
+            plaintext_ballot, internal_manifest, context, compact_ballot.ballot_nonce
+        )
+    )
+    crypto_hash = create_ballot_hash(
+        plaintext_ballot.object_id, internal_manifest.manifest_hash, contests
+    )
+
+    return SubmittedBallot(
+        plaintext_ballot.object_id,
+        plaintext_ballot.style_id,
+        internal_manifest.manifest_hash,
+        compact_ballot.previous_code,
+        contests,
+        compact_ballot.code,
+        compact_ballot.timestamp,
+        crypto_hash,
+        compact_ballot.ballot_nonce,
+        compact_ballot.ballot_box_state,
+    )
+
+
+def expand_compact_plaintext_ballot(
+    compact_ballot: CompactPlaintextBallot, internal_manifest: InternalManifest
+) -> PlaintextBallot:
+    """Expand a compact plaintext ballot into the original plaintext ballot"""
+    return PlaintextBallot(
+        compact_ballot.object_id,
+        compact_ballot.style_id,
+        _get_plaintext_contests(compact_ballot, internal_manifest),
+    )
+
+
+def _get_compact_selections(ballot: PlaintextBallot) -> List[bool]:
+    selections = []
+    for contest in ballot.contests:
+        for selection in contest.ballot_selections:
+            selections.append(selection.vote == YES_VOTE)
+    return selections
+
+
+def _get_compact_extended_data(ballot: PlaintextBallot) -> Dict[int, ExtendedData]:
+    extended_data = {}
+    index = 0
+    for contest in ballot.contests:
+        for selection in contest.ballot_selections:
+            index += 1
+            if selection.extended_data:
+                extended_data[index] = selection.extended_data
+    return extended_data
+
+
+def _get_plaintext_contests(
+    compact_ballot: CompactPlaintextBallot, internal_manifest: InternalManifest
+) -> List[PlaintextBallotContest]:
+    """Get ballot contests from compact plaintext ballot"""
+    index = 0
+    ballot_style_contests = _get_ballot_style_contests(
+        compact_ballot.style_id, internal_manifest
+    )
+
+    contests: List[PlaintextBallotContest] = []
+    for manifest_contest in sorted(
+        internal_manifest.contests, key=lambda c: c.sequence_order
+    ):
+        contest_in_style = (
+            ballot_style_contests.get(manifest_contest.object_id) is not None
+        )
+
+        # Iterate through selections. If contest not in style, mark placeholder
+        selections: List[PlaintextBallotSelection] = []
+        for selection in sorted(
+            manifest_contest.ballot_selections, key=lambda s: s.sequence_order
+        ):
+            selections.append(
+                PlaintextBallotSelection(
+                    selection.object_id,
+                    YES_VOTE if compact_ballot.selections[index] else NO_VOTE,
+                    not contest_in_style,
+                    compact_ballot.extended_data.get(index),
+                )
+            )
+            index += 1
+
+        contests.append(PlaintextBallotContest(manifest_contest.object_id, selections))
+    return contests
+
+
+def _get_ballot_style_contests(
+    ballot_style_id: str, internal_manifest: InternalManifest
+) -> Dict[str, ContestDescriptionWithPlaceholders]:
+    ballot_style_contests = internal_manifest.get_contests_for(ballot_style_id)
+    return {contest.object_id: contest for contest in ballot_style_contests}

src/electionguard/encrypt.py

@@ -35,18 +35,40 @@ class EncryptionDevice(Serializable):
     """
 
     uuid: int
+    """Unique identifier for device"""
+
+    session_id: str
+    """Used to identify session and protect the timestamp"""
+
+    launch_code: int
+    """Election initialization value"""
+
     location: str
+    """Arbitary string to designate the location of device"""
 
-    def __init__(self, location: str) -> None:
-        self.uuid = generate_device_uuid()
+    def __init__(
+        self,
+        uuid: int,
+        session_id: str,
+        launch_code: int,
+        location: str,
+    ) -> None:
+        self.uuid = uuid
+        self.session_id = session_id
+        self.launch_code = launch_code
         self.location = location
 
     def get_hash(self) -> ElementModQ:
         """
         Get hash for encryption device
         :return: Starting hash
         """
-        return get_hash_for_device(self.uuid, self.location)
+        return get_hash_for_device(
+            self.uuid, self.session_id, self.launch_code, self.location
+        )
+
+    def get_timestamp(self) -> int:
+        pass
 
 
 class EncryptionMediator:
@@ -420,30 +442,11 @@ def encrypt_ballot(
         random_master_nonce,
     )
 
-    encrypted_contests: List[CiphertextBallotContest] = list()
-
-    # only iterate on contests for this specific ballot style
-    for description in internal_manifest.get_contests_for(ballot.style_id):
-        use_contest = None
-        for contest in ballot.contests:
-            if contest.object_id == description.object_id:
-                use_contest = contest
-                break
-        # no selections provided for the contest, so create a placeholder contest
-        if not use_contest:
-            use_contest = contest_from(description)
-
-        encrypted_contest = encrypt_contest(
-            use_contest,
-            description,
-            context.elgamal_public_key,
-            context.crypto_extended_base_hash,
-            nonce_seed,
-        )
-
-        if encrypted_contest is None:
-            return None  # log will have happened earlier
-        encrypted_contests.append(get_optional(encrypted_contest))
+    encrypted_contests = encrypt_ballot_contests(
+        ballot, internal_manifest, context, nonce_seed
+    )
+    if encrypted_contests is None:
+        return None
 
     # Create the return object
     encrypted_ballot = make_ciphertext_ballot(
@@ -469,3 +472,38 @@ def encrypt_ballot(
     ):
         return encrypted_ballot
     return None  # log will have happened earlier
+
+
+def encrypt_ballot_contests(
+    ballot: PlaintextBallot,
+    description: InternalManifest,
+    context: CiphertextElectionContext,
+    nonce_seed: ElementModQ,
+) -> Optional[List[CiphertextBallotContest]]:
+    """Encrypt contests from a plaintext ballot with a specific style"""
+    encrypted_contests: List[CiphertextBallotContest] = []
+
+    # Only iterate on contests for this specific ballot style
+    for ballot_style_contest in description.get_contests_for(ballot.style_id):
+        use_contest = None
+        for contest in ballot.contests:
+            if contest.object_id == ballot_style_contest.object_id:
+                use_contest = contest
+                break
+
+        # no selections provided for the contest, so create a placeholder contest
+        if not use_contest:
+            use_contest = contest_from(ballot_style_contest)
+
+        encrypted_contest = encrypt_contest(
+            use_contest,
+            ballot_style_contest,
+            context.elgamal_public_key,
+            context.crypto_extended_base_hash,
+            nonce_seed,
+        )
+
+        if encrypted_contest is None:
+            return None
+        encrypted_contests.append(get_optional(encrypted_contest))
+    return encrypted_contests

src/electionguardtest/election_factory.py

@@ -1,5 +1,6 @@
 from datetime import datetime
 import os
+import uuid
 from dataclasses import dataclass
 from typing import TypeVar, Callable, Optional, Tuple, List
 
@@ -14,7 +15,7 @@
 from electionguard.ballot import PlaintextBallot
 from electionguard.election import CiphertextElectionContext, ElectionConstants
 from electionguard.election_builder import ElectionBuilder
-from electionguard.encrypt import contest_from
+from electionguard.encrypt import EncryptionDevice, contest_from, generate_device_uuid
 from electionguard.group import ElementModP, TWO_MOD_Q
 from electionguard.guardian import Guardian
 from electionguard.key_ceremony import CoefficientValidationSet
@@ -257,6 +258,15 @@ def _get_manifest_from_file(filename: str) -> Manifest:
 
         return manifest
 
+    @staticmethod
+    def get_encryption_device() -> EncryptionDevice:
+        return EncryptionDevice(
+            generate_device_uuid(),
+            "Session",
+            12345,
+            f"polling-place-{str(uuid.uuid1())}",
+        )
+
 
 @composite
 def get_selection_description_well_formed(