Allow specifying sample to use when running the generator (#511)

* Allow specifying sample to use when running the generator

* try unzip sample data

* fix spacing

* run lint

.gitignore

@@ -6,6 +6,8 @@ __pycache__/
 # C extensions
 *.so
 
+data/0.95.0/
+
 # Distribution / packaging
 .Python
 build/
@@ -145,4 +147,7 @@ election_record.zip
 election_private_data.zip
 
 # VS Code
-.vscode/settings.json
+.vscode/settings.json
+sample-data.zip
+
+results/

Makefile

@@ -1,4 +1,4 @@
-.PHONY: all openssl-fix install install-gmp install-gmp-mac install-gmp-linux install-gmp-windows install-mkdocs auto-lint validate test test-example bench coverage coverage-html coverage-xml coverage-erase generate-sample-data
+.PHONY: all environment openssl-fix install install-gmp install-gmp-mac install-gmp-linux install-gmp-windows install-mkdocs auto-lint validate test test-example bench coverage coverage-html coverage-xml coverage-erase generate-sample-data
 
 CODE_COVERAGE ?= 90
 OS ?= $(shell python -c 'import platform; print(platform.system())')
@@ -19,6 +19,13 @@ environment:
 	poetry config virtualenvs.in-project true 
 	poetry install
 	@echo 🚨 Be sure to add poetry to PATH
+ifeq ($(OS), Windows)
+	choco install wget
+	choco install unzip
+endif
+	wget https://github.com/microsoft/electionguard/releases/download/v0.95.0/sample-data.zip
+	unzip -o sample-data.zip
+	unzip sample-data.zip
 
 install:
 	@echo 🔧 INSTALL
@@ -181,7 +188,7 @@ dependency-graph-ci:
 
 # Sample Data
 generate-sample-data:
-	poetry run python3 src/electionguard_tools/scripts/sample_generator.py -n $(SAMPLE_BALLOT_COUNT) -s $(SAMPLE_BALLOT_SPOIL_RATE)
+	poetry run python3 src/electionguard_tools/scripts/sample_generator.py -m "hamilton-general" -n $(SAMPLE_BALLOT_COUNT) -s $(SAMPLE_BALLOT_SPOIL_RATE)
 
 # Publish
 publish:

src/electionguard_tools/__init__.py

@@ -68,6 +68,8 @@ from electionguard_tools.helpers import (
 )
 from electionguard_tools.scripts import (
     DEFAULT_NUMBER_OF_BALLOTS,
+    DEFAULT_SAMPLE_MANIFEST,
+    DEFAULT_SPEC_VERSION,
     DEFAULT_SPOIL_RATE,
     DEFAULT_USE_ALL_GUARDIANS,
     DEFAULT_USE_PRIVATE_DATA,
@@ -121,6 +123,8 @@ __all__ = [
     "CONTEXT_FILE_NAME",
     "CiphertextElectionsTupleType",
     "DEFAULT_NUMBER_OF_BALLOTS",
+    "DEFAULT_SAMPLE_MANIFEST",
+    "DEFAULT_SPEC_VERSION",
     "DEFAULT_SPOIL_RATE",
     "DEFAULT_USE_ALL_GUARDIANS",
     "DEFAULT_USE_PRIVATE_DATA",

src/electionguard_tools/factories/election_factory.py

@@ -83,6 +83,16 @@ class ElectionFactory:
         """Get simple manifest from json file."""
         return self._get_manifest_from_file(self.simple_election_manifest_file_name)
 
+    @staticmethod
+    def get_manifest_from_file(spec_version: str, sample_manifest: str) -> Manifest:
+        """Get simple manifest from json file."""
+        return from_file_to_dataclass(
+            Manifest,
+            os.path.join(
+                data, spec_version, "sample", sample_manifest, "manifest.json"
+            ),
+        )
+
     @staticmethod
     def get_hamilton_manifest_from_file() -> Manifest:
         """Get Hamilton County manifest from json file."""
@@ -93,15 +103,15 @@ class ElectionFactory:
             ),
         )
 
-    def get_hamilton_manifest_with_encryption_context(
-        self,
+    def get_sample_manifest_with_encryption_context(
+        self, spec_version: str, sample_manifest: str
     ) -> Tuple[AllPublicElectionData, AllPrivateElectionData]:
         """Get hamilton manifest and context"""
         guardians: List[Guardian] = []
         guardian_records: List[GuardianRecord] = []
 
         # Configure the election builder
-        manifest = self.get_hamilton_manifest_from_file()
+        manifest = self.get_manifest_from_file(spec_version, sample_manifest)
         builder = ElectionBuilder(NUMBER_OF_GUARDIANS, QUORUM, manifest)
 
         # Run the Key Ceremony

src/electionguard_tools/scripts/__init__.py

@@ -2,6 +2,8 @@ from electionguard_tools.scripts import sample_generator
 
 from electionguard_tools.scripts.sample_generator import (
     DEFAULT_NUMBER_OF_BALLOTS,
+    DEFAULT_SAMPLE_MANIFEST,
+    DEFAULT_SPEC_VERSION,
     DEFAULT_SPOIL_RATE,
     DEFAULT_USE_ALL_GUARDIANS,
     DEFAULT_USE_PRIVATE_DATA,
@@ -10,6 +12,8 @@ from electionguard_tools.scripts.sample_generator import (
 
 __all__ = [
     "DEFAULT_NUMBER_OF_BALLOTS",
+    "DEFAULT_SAMPLE_MANIFEST",
+    "DEFAULT_SPEC_VERSION",
     "DEFAULT_SPOIL_RATE",
     "DEFAULT_USE_ALL_GUARDIANS",
     "DEFAULT_USE_PRIVATE_DATA",

src/electionguard_tools/scripts/sample_generator.py

@@ -11,6 +11,7 @@ from electionguard.ballot import (
 from electionguard.data_store import DataStore
 from electionguard.ballot_box import BallotBox, get_ballots
 from electionguard.decryption_mediator import DecryptionMediator
+from electionguard.election_polynomial import LagrangeCoefficientsRecord
 from electionguard.encrypt import (
     EncryptionDevice,
     EncryptionMediator,
@@ -36,6 +37,8 @@ DEFAULT_NUMBER_OF_BALLOTS = 5
 DEFAULT_SPOIL_RATE = 50
 DEFAULT_USE_ALL_GUARDIANS = False
 DEFAULT_USE_PRIVATE_DATA = False
+DEFAULT_SPEC_VERSION = "0.95.0"
+DEFAULT_SAMPLE_MANIFEST = "hamilton-general"
 
 
 class ElectionSampleDataGenerator:
@@ -61,6 +64,8 @@ class ElectionSampleDataGenerator:
         spoil_rate: int = DEFAULT_SPOIL_RATE,
         use_all_guardians: bool = DEFAULT_USE_ALL_GUARDIANS,
         use_private_data: bool = DEFAULT_USE_PRIVATE_DATA,
+        spec_version: str = DEFAULT_SPEC_VERSION,
+        sample_manifest: str = DEFAULT_SAMPLE_MANIFEST,
     ):
         """
         Generate the sample data set
@@ -71,10 +76,13 @@ class ElectionSampleDataGenerator:
         rmtree(PRIVATE_DATA_DIR, ignore_errors=True)
 
         # Configure the election
+        # TODO: pass the spec version and the manifest name in
         (
             manifest,
             private_data,
-        ) = self.election_factory.get_hamilton_manifest_with_encryption_context()
+        ) = self.election_factory.get_sample_manifest_with_encryption_context(
+            spec_version, sample_manifest
+        )
         plaintext_ballots = (
             self.ballot_factory.generate_fake_plaintext_ballots_for_election(
                 manifest.internal_manifest, number_of_ballots
@@ -158,6 +166,9 @@ class ElectionSampleDataGenerator:
                 ciphertext_tally.publish(),
                 plaintext_tally,
                 manifest.guardians,
+                LagrangeCoefficientsRecord(
+                    list(mediator.get_lagrange_coefficients().values())
+                ),
             )
 
             if use_private_data:
@@ -188,6 +199,16 @@ if __name__ == "__main__":
         description="Generate sample ballot data",
         formatter_class=argparse.ArgumentDefaultsHelpFormatter,
     )
+
+    parser.add_argument(
+        "-m",
+        "--manifest",
+        metavar="<manifest>",
+        default=DEFAULT_SAMPLE_MANIFEST,
+        type=str,
+        help="The manifest to use to generate sample data.",
+        choices=["full", "hamilton-general", "minimal", "small"],
+    )
     parser.add_argument(
         "-n",
         "--number-of-ballots",
@@ -218,8 +239,21 @@ if __name__ == "__main__":
         action="store_true",
         help="Include private data when generating.",
     )
+    parser.add_argument(
+        "-v",
+        "--version",
+        default=DEFAULT_SPEC_VERSION,
+        type=str,
+        help="The spec version to use.",
+        choices=[DEFAULT_SPEC_VERSION],
+    )
     args = parser.parse_args()
 
     ElectionSampleDataGenerator().generate(
-        args.number_of_ballots, args.spoil_rate, args.all_guardians, args.private_data
+        args.number_of_ballots,
+        args.spoil_rate,
+        args.all_guardians,
+        args.private_data,
+        args.version,
+        args.manifest,
     )