#776 USB Import Wizard (#781)

* Better icon

* Refactor existing upload ballot logic into a sub-component

* Show wizard if we're on windows else legacy

* Ability to scan for drives and display results

* can close the wizard and select manually

* Ability to upload ballots, no UI feedback yet

* refactor common success code, implement upload more

* add scan button

* Show status of uploading ballots

* Duplicate handeling

* Add cancel button to create decryption

* Show status on decryption page

* fix linting issue

* Show status updates while guardians are decrypting

* better unhappy path notifications

* Poll for usb devices

Author: lprichar

src/electionguard_gui/__init__.py

@@ -33,6 +33,7 @@
     key_ceremony_details_component,
     notify_ui_db_changed,
     refresh_decryption,
+    update_upload_status,
     upload_ballots_component,
     view_decryption_component,
     view_election_component,
@@ -112,12 +113,12 @@
     election_service,
     export_service,
     get_data_dir,
-    get_drives,
     get_export_dir,
     get_export_locations,
     get_guardian_number,
     get_key_ceremony_status,
     get_plaintext_ballot_report,
+    get_removable_drives,
     get_tally,
     guardian_service,
     gui_setup_input_retrieval_step,
@@ -139,6 +140,7 @@
     service_base,
     status_descriptions,
     to_ballot_share_raw,
+    update_decrypt_status,
     verification_to_dict,
     version_service,
 )
@@ -232,12 +234,12 @@
     "export_encryption_package_component",
     "export_service",
     "get_data_dir",
-    "get_drives",
     "get_export_dir",
     "get_export_locations",
     "get_guardian_number",
     "get_key_ceremony_status",
     "get_plaintext_ballot_report",
+    "get_removable_drives",
     "get_spoiled_ballot_by_id",
     "get_tally",
     "guardian_home_component",
@@ -271,6 +273,8 @@
     "start",
     "status_descriptions",
     "to_ballot_share_raw",
+    "update_decrypt_status",
+    "update_upload_status",
     "upload_ballots_component",
     "utc_to_str",
     "verification_to_dict",

src/electionguard_gui/components/__init__.py

@@ -43,6 +43,7 @@
 )
 from electionguard_gui.components.upload_ballots_component import (
     UploadBallotsComponent,
+    update_upload_status,
 )
 from electionguard_gui.components.view_decryption_component import (
     ViewDecryptionComponent,
@@ -86,6 +87,7 @@
     "key_ceremony_details_component",
     "notify_ui_db_changed",
     "refresh_decryption",
+    "update_upload_status",
     "upload_ballots_component",
     "view_decryption_component",
     "view_election_component",

src/electionguard_gui/components/upload_ballots_component.py

@@ -1,11 +1,14 @@
+import os
 from typing import Any
 from datetime import datetime
 import eel
-from electionguard.serialize import from_raw
+from electionguard.encrypt import EncryptionDevice
+from electionguard.serialize import from_file, from_raw
 from electionguard.ballot import SubmittedBallot
 from electionguard_gui.components.component_base import ComponentBase
 from electionguard_gui.eel_utils import eel_fail, eel_success
 from electionguard_gui.services import ElectionService, BallotUploadService
+from electionguard_gui.services.export_service import get_removable_drives
 
 
 class UploadBallotsComponent(ComponentBase):
@@ -25,6 +28,9 @@ def __init__(
     def expose(self) -> None:
         eel.expose(self.create_ballot_upload)
         eel.expose(self.upload_ballot)
+        eel.expose(self.is_wizard_supported)
+        eel.expose(self.scan_drives)
+        eel.expose(self.upload_ballots)
 
     def create_ballot_upload(
         self,
@@ -106,3 +112,116 @@ def upload_ballot(
         # pylint: disable=broad-except
         except Exception as e:
             return self.handle_error(e)
+
+    # pylint: disable=no-self-use
+    def is_wizard_supported(self) -> bool:
+        on_windows = os.name == "nt"
+        return on_windows
+
+    def scan_drives(self) -> dict[str, Any]:
+        try:
+            removable_drives = get_removable_drives()
+            self._log.trace(f"found {len(removable_drives)} removable drives")
+            candidate_drives = [
+                self.parse_drive(drive)
+                for drive in removable_drives
+                if os.path.exists(os.path.join(drive, "artifacts", "encrypted_ballots"))
+                and os.path.exists(os.path.join(drive, "artifacts", "devices"))
+            ]
+            first_candidate = next(iter(candidate_drives), None)
+            return eel_success(first_candidate)
+        # pylint: disable=broad-except
+        except Exception as e:
+            return self.handle_error(e)
+
+    def parse_drive(self, drive: str) -> dict[str, Any]:
+        ballots_dir = os.path.join(drive, "artifacts", "encrypted_ballots")
+        devices_dir = os.path.join(drive, "artifacts", "devices")
+        device_files = os.listdir(devices_dir)
+        device_file_name = next(iter(os.listdir(devices_dir)))
+        device_file_path = os.path.join(devices_dir, device_file_name)
+        if len(device_files) > 1:
+            self._log.warn(
+                "found multiple device files in drive, using " + device_file_name
+            )
+        device_file_json = from_file(EncryptionDevice, device_file_path)
+        location = device_file_json.location
+        ballot_count = len(os.listdir(ballots_dir))
+        return {
+            "drive": drive,
+            "ballots": ballot_count,
+            "location": location,
+            "device_file_name": device_file_name,
+            "device_file_path": device_file_path,
+            "ballots_dir": ballots_dir,
+        }
+
+    def upload_ballots(self, election_id: str) -> dict[str, Any]:
+        try:
+            update_upload_status("Scanning drives")
+            drive_info = self.scan_drives()
+            device_file_name = drive_info["result"]["device_file_name"]
+            device_file_path = drive_info["result"]["device_file_path"]
+            self._log.debug(
+                f"uploading ballots for {election_id} from {device_file_path} device {device_file_name}"
+            )
+            update_upload_status("Uploading device file")
+            ballot_upload_result = self.create_ballot_upload_from_file(
+                election_id,
+                device_file_name,
+                device_file_path,
+            )
+            if not ballot_upload_result["success"]:
+                return ballot_upload_result
+
+            ballots_dir: str = drive_info["result"]["ballots_dir"]
+            ballot_files = os.listdir(ballots_dir)
+            ballot_upload_id: str = ballot_upload_result["result"]
+            ballot_num = 1
+            duplicate_count = 0
+            ballot_count = len(ballot_files)
+            for ballot_file in ballot_files:
+                self._log.debug("uploading ballot " + ballot_file)
+                update_upload_status(f"Uploading ballot {ballot_num}/{ballot_count}")
+                result = self.create_ballot_from_file(
+                    election_id, ballot_file, ballot_upload_id, ballots_dir
+                )
+                if not result["success"]:
+                    return result
+                if result["result"]["is_duplicate"]:
+                    duplicate_count += 1
+                ballot_num += 1
+            return eel_success(
+                {"ballot_count": ballot_count, "duplicate_count": duplicate_count}
+            )
+        # pylint: disable=broad-except
+        except Exception as e:
+            return self.handle_error(e)
+
+    def create_ballot_from_file(
+        self,
+        election_id: str,
+        ballot_file_name: str,
+        ballot_upload_id: str,
+        ballots_dir: str,
+    ) -> dict[str, Any]:
+        ballot_file_path = os.path.join(ballots_dir, ballot_file_name)
+        with open(ballot_file_path, "r", encoding="utf-8") as ballot_file:
+            ballot_contents = ballot_file.read()
+            return self.upload_ballot(
+                ballot_upload_id, election_id, ballot_file_name, ballot_contents
+            )
+
+    def create_ballot_upload_from_file(
+        self, election_id: str, device_file_name: str, device_file_path: str
+    ) -> dict[str, Any]:
+        with open(device_file_path, "r", encoding="utf-8") as device_file:
+            ballot_upload = self.create_ballot_upload(
+                election_id, device_file_name, device_file.read()
+            )
+            return ballot_upload
+
+
+def update_upload_status(status: str) -> None:
+    # pylint: disable=no-member
+    eel.update_upload_status(status)

src/electionguard_gui/services/__init__.py

@@ -58,6 +58,7 @@
     decryption_s2_announce_service,
     decryption_stage_base,
     get_tally,
+    update_decrypt_status,
 )
 from electionguard_gui.services.directory_service import (
     DOCKER_MOUNT_DIR,
@@ -71,8 +72,8 @@
     ElectionService,
 )
 from electionguard_gui.services.export_service import (
-    get_drives,
     get_export_locations,
+    get_removable_drives,
 )
 from electionguard_gui.services.guardian_service import (
     GuardianService,
@@ -168,12 +169,12 @@
     "election_service",
     "export_service",
     "get_data_dir",
-    "get_drives",
     "get_export_dir",
     "get_export_locations",
     "get_guardian_number",
     "get_key_ceremony_status",
     "get_plaintext_ballot_report",
+    "get_removable_drives",
     "get_tally",
     "guardian_service",
     "gui_setup_input_retrieval_step",
@@ -195,6 +196,7 @@
     "service_base",
     "status_descriptions",
     "to_ballot_share_raw",
+    "update_decrypt_status",
     "verification_to_dict",
     "version_service",
 ]

src/electionguard_gui/services/decryption_stages/__init__.py

@@ -7,6 +7,7 @@
 )
 from electionguard_gui.services.decryption_stages.decryption_s2_announce_service import (
     DecryptionS2AnnounceService,
+    update_decrypt_status,
 )
 from electionguard_gui.services.decryption_stages.decryption_stage_base import (
     DecryptionStageBase,
@@ -21,4 +22,5 @@
     "decryption_s2_announce_service",
     "decryption_stage_base",
     "get_tally",
+    "update_decrypt_status",
 ]

src/electionguard_gui/services/decryption_stages/decryption_s1_join_service.py

@@ -1,4 +1,5 @@
 from pymongo.database import Database
+import eel
 from electionguard.ballot import BallotBoxState
 
 from electionguard_gui.models.decryption_dto import DecryptionDto
@@ -12,6 +13,7 @@ class DecryptionS1JoinService(DecryptionStageBase):
     """Responsible for the 1st stage during a decryption were guardians join the decryption"""
 
     def run(self, db: Database, decryption: DecryptionDto) -> None:
+        update_decrypt_status("Starting tally")
         current_user_id = self._auth_service.get_required_user_id()
         self._log.info(f"S1: {current_user_id} decrypting  {decryption.decryption_id}")
         election = self._election_service.get(db, decryption.election_id)
@@ -22,18 +24,22 @@ def run(self, db: Database, decryption: DecryptionDto) -> None:
             current_user_id, decryption
         )
         ballots = self._ballot_upload_service.get_ballots(db, election.id)
-        spoiled_ballots = [
-            ballot for ballot in ballots if ballot.state == BallotBoxState.SPOILED
-        ]
+        update_decrypt_status("Calculating tally")
         ciphertext_tally = get_tally(manifest, context, ballots, False)
         decryption_share = guardian.compute_tally_share(ciphertext_tally, context)
         if decryption_share is None:
             raise Exception("No decryption_shares found")
+
+        update_decrypt_status("Calculating spoiled ballots")
+        spoiled_ballots = [
+            ballot for ballot in ballots if ballot.state == BallotBoxState.SPOILED
+        ]
         ballot_shares = guardian.compute_ballot_shares(spoiled_ballots, context)
         if ballot_shares is None:
             raise Exception("No ballot shares found")
         guardian_key = guardian.share_key()
 
+        update_decrypt_status("Finalizing tally")
         self._decryption_service.append_guardian_joined(
             db,
             decryption.decryption_id,
@@ -43,3 +49,8 @@ def run(self, db: Database, decryption: DecryptionDto) -> None:
             guardian_key,
         )
         self._decryption_service.notify_changed(db, decryption.decryption_id)
+
+
+def update_decrypt_status(status: str) -> None:
+    # pylint: disable=no-member
+    eel.update_decrypt_status(status)

src/electionguard_gui/services/decryption_stages/decryption_s2_announce_service.py

@@ -1,4 +1,5 @@
 from pymongo.database import Database
+import eel
 from electionguard import DecryptionMediator
 from electionguard.ballot import BallotBoxState
 from electionguard.election_polynomial import LagrangeCoefficientsRecord
@@ -19,6 +20,7 @@ def should_run(self, db: Database, decryption: DecryptionDto) -> bool:
         return is_admin and all_guardians_joined and not is_completed
 
     def run(self, db: Database, decryption: DecryptionDto) -> None:
+        update_decrypt_status("Starting tally")
         self._log.info(f"S2: Announcing decryption {decryption.decryption_id}")
         election = self._election_service.get(db, decryption.election_id)
         context = election.get_context()
@@ -28,7 +30,10 @@ def run(self, db: Database, decryption: DecryptionDto) -> None:
             context,
         )
         decryption_shares = decryption.get_decryption_shares()
+        share_count = len(decryption_shares)
+        current_share = 1
         for decryption_share_dict in decryption_shares:
+            update_decrypt_status(f"Calculating share {current_share}/{share_count}")
             self._log.debug(f"announcing {decryption_share_dict.guardian_id}")
             guardian_sequence_number = election.get_guardian_sequence_order(
                 decryption_share_dict.guardian_id
@@ -41,7 +46,9 @@ def run(self, db: Database, decryption: DecryptionDto) -> None:
                 decryption_share_dict.tally_share,
                 decryption_share_dict.ballot_shares,
             )
+            current_share += 1
 
+        update_decrypt_status("Decrypting spoiled ballots")
         manifest = election.get_manifest()
         ballots = self._ballot_upload_service.get_ballots(db, election.id)
         spoiled_ballots = [
@@ -61,6 +68,8 @@ def run(self, db: Database, decryption: DecryptionDto) -> None:
         if plaintext_spoiled_ballots is None:
             raise Exception("No plaintext spoiled ballots found")
 
+        update_decrypt_status("Finalizing tally")
+
         lagrange_coefficients = _get_lagrange_coefficients(decryption_mediator)
 
         self._log.debug("setting decryption completed")
@@ -80,3 +89,8 @@ def _get_lagrange_coefficients(
     decryption_mediator: DecryptionMediator,
 ) -> LagrangeCoefficientsRecord:
     return LagrangeCoefficientsRecord(decryption_mediator.get_lagrange_coefficients())
+
+
+def update_decrypt_status(status: str) -> None:
+    # pylint: disable=no-member
+    eel.update_decrypt_status(status)

src/electionguard_gui/services/export_service.py

@@ -5,13 +5,13 @@
 def get_export_locations() -> list[str]:
     export_dir = get_export_dir()
     if os.name == "nt":
-        drives = get_drives()
+        drives = get_removable_drives()
         return [export_dir, _get_download_path(), get_data_dir()] + drives
     return [export_dir]
 
 
-def get_drives() -> list[str]:
-    dl = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"
+def get_removable_drives() -> list[str]:
+    dl = "DEFGHIJKLMNOPQRSTUVWXYZ"
     drives = [f"{d}:\\" for d in dl if os.path.exists(f"{d}:")]
     return drives