🔓 Ballot decryption endpoints for all guardians (#94)

* 🔓 Ballot decryption endpoints

* Change ballot share route to `decrypt-shares` for clarity

* Pylint fixes

* Move ballot indexing to a helper function

Author: rkorsak

app/api/v1/guardian/ballot.py

@@ -0,0 +1,41 @@
+from typing import Any
+from electionguard.ballot import CiphertextAcceptedBallot
+from electionguard.decryption import compute_decryption_share_for_ballot
+from electionguard.election import CiphertextElectionContext
+from electionguard.scheduler import Scheduler
+from electionguard.serializable import write_json_object
+from fastapi import APIRouter, Body
+
+from ..models import (
+    convert_guardian,
+    DecryptBallotSharesRequest,
+    DecryptBallotSharesResponse,
+)
+from ..tags import TALLY
+
+router = APIRouter()
+
+
+@router.post("/decrypt-shares", tags=[TALLY])
+def decrypt_ballot_shares(request: DecryptBallotSharesRequest = Body(...)) -> Any:
+    """
+    Decrypt this guardian's share of one or more ballots
+    """
+    ballots = [
+        CiphertextAcceptedBallot.from_json_object(ballot)
+        for ballot in request.encrypted_ballots
+    ]
+    context = CiphertextElectionContext.from_json_object(request.context)
+    guardian = convert_guardian(request.guardian)
+
+    scheduler = Scheduler()
+    shares = [
+        compute_decryption_share_for_ballot(guardian, ballot, context, scheduler)
+        for ballot in ballots
+    ]
+
+    response = DecryptBallotSharesResponse(
+        shares=[write_json_object(share) for share in shares]
+    )
+
+    return response

app/api/v1/guardian/routes.py

@@ -1,4 +1,5 @@
 from fastapi import APIRouter
+from . import ballot
 from . import guardian
 from . import key
 from . import tally
@@ -7,4 +8,5 @@
 
 router.include_router(guardian.router, prefix="/guardian")
 router.include_router(key.router, prefix="/key")
+router.include_router(ballot.router, prefix="/ballot")
 router.include_router(tally.router, prefix="/tally")

app/api/v1/mediator/ballot.py

@@ -1,6 +1,12 @@
-from typing import Any, Optional
-from electionguard.ballot import CiphertextBallot, PlaintextBallot
+from typing import Any, Dict, List, Optional
+from electionguard.ballot import (
+    CiphertextAcceptedBallot,
+    CiphertextBallot,
+    PlaintextBallot,
+)
 from electionguard.ballot_box import accept_ballot, BallotBoxState
+from electionguard.decrypt_with_shares import decrypt_ballot
+from electionguard.decryption_share import BallotDecryptionShare
 from electionguard.ballot_store import BallotStore
 from electionguard.election import (
     CiphertextElectionContext,
@@ -10,11 +16,17 @@
 from electionguard.encrypt import encrypt_ballot
 from electionguard.group import ElementModQ
 from electionguard.serializable import read_json_object, write_json_object
+from electionguard.types import BALLOT_ID, GUARDIAN_ID
 from electionguard.utils import get_optional
 from fastapi import APIRouter, Body, HTTPException
 
-from ..models import AcceptBallotRequest, EncryptBallotsRequest, EncryptBallotsResponse
-from ..tags import CAST_AND_SPOIL, ENCRYPT_BALLOTS
+from ..models import (
+    AcceptBallotRequest,
+    DecryptBallotsRequest,
+    EncryptBallotsRequest,
+    EncryptBallotsResponse,
+)
+from ..tags import CAST_AND_SPOIL, ENCRYPT_BALLOTS, TALLY
 
 router = APIRouter()
 
@@ -33,6 +45,34 @@ def cast_ballot(request: AcceptBallotRequest = Body(...)) -> Any:
     return casted_ballot.to_json_object()
 
 
+@router.post("/decrypt", tags=[TALLY])
+def decrypt_ballots(request: DecryptBallotsRequest = Body(...)) -> Any:
+    ballots = [
+        CiphertextAcceptedBallot.from_json_object(ballot)
+        for ballot in request.encrypted_ballots
+    ]
+    context: CiphertextElectionContext = CiphertextElectionContext.from_json_object(
+        request.context
+    )
+
+    all_shares: List[BallotDecryptionShare] = [
+        read_json_object(share, BallotDecryptionShare)
+        for shares in request.shares.values()
+        for share in shares
+    ]
+    shares_by_ballot = index_shares_by_ballot(all_shares)
+
+    extended_base_hash = context.crypto_extended_base_hash
+    decrypted_ballots = {
+        ballot.object_id: decrypt_ballot(
+            ballot, shares_by_ballot[ballot.object_id], extended_base_hash
+        )
+        for ballot in ballots
+    }
+
+    return write_json_object(decrypted_ballots)
+
+
 @router.post("/spoil", tags=[CAST_AND_SPOIL])
 def spoil_ballot(request: AcceptBallotRequest = Body(...)) -> Any:
     """
@@ -96,3 +136,18 @@ def handle_ballot(request: AcceptBallotRequest, state: BallotBoxState) -> Any:
     )
 
     return accepted_ballot
+
+
+def index_shares_by_ballot(
+    shares: List[BallotDecryptionShare],
+) -> Dict[BALLOT_ID, Dict[GUARDIAN_ID, BallotDecryptionShare]]:
+    """
+    Construct a lookup by ballot ID containing the dictionary of shares needed
+    to decrypt that ballot.
+    """
+    shares_by_ballot: Dict[str, Dict[str, BallotDecryptionShare]] = {}
+    for share in shares:
+        ballot_shares = shares_by_ballot.setdefault(share.ballot_id, {})
+        ballot_shares[share.guardian_id] = share
+
+    return shares_by_ballot

app/api/v1/mediator/tally.py

@@ -12,7 +12,6 @@
     publish_ciphertext_tally,
     publish_plaintext_tally,
     CiphertextTally,
-    PublishedCiphertextTally,
 )
 from fastapi import APIRouter, Body, HTTPException
 
@@ -70,6 +69,11 @@ def decrypt_tally(request: DecryptTallyRequest = Body(...)) -> Any:
     }
 
     full_plaintext_tally = decrypt(tally, shares, context)
+    if not full_plaintext_tally:
+        raise HTTPException(
+            status_code=500,
+            detail="Unable to decrypt tally",
+        )
     published_plaintext_tally = publish_plaintext_tally(full_plaintext_tally)
 
     return published_plaintext_tally.to_json_object()
@@ -97,7 +101,7 @@ def _parse_tally_request(
 
 def _tally_ballots(
     tally: CiphertextTally, ballots: List[CiphertextAcceptedBallot]
-) -> PublishedCiphertextTally:
+) -> Any:
     """
     Append a series of ballots to a new or existing tally
     """

app/api/v1/models/ballot.py

@@ -1,18 +1,23 @@
-from typing import Any, List, Optional
+from typing import Any, Dict, List, Optional
 
 from .base import Base
 from .election import ElectionDescription, CiphertextElectionContext
+from .guardian import Guardian, GuardianId
 
 
 __all__ = [
+    "AcceptBallotRequest",
     "CiphertextAcceptedBallot",
     "CiphertextBallot",
-    "PlaintextBallot",
-    "AcceptBallotRequest",
+    "DecryptBallotSharesRequest",
+    "DecryptBallotSharesResponse",
+    "DecryptBallotsRequest",
     "EncryptBallotsRequest",
     "EncryptBallotsResponse",
+    "PlaintextBallot",
 ]
 
+BallotDecryptionShare = Any
 CiphertextAcceptedBallot = Any
 CiphertextBallot = Any
 PlaintextBallot = Any
@@ -24,6 +29,22 @@ class AcceptBallotRequest(Base):
     context: CiphertextElectionContext
 
 
+class DecryptBallotsRequest(Base):
+    encrypted_ballots: List[CiphertextAcceptedBallot]
+    shares: Dict[GuardianId, List[BallotDecryptionShare]]
+    context: CiphertextElectionContext
+
+
+class DecryptBallotSharesRequest(Base):
+    encrypted_ballots: List[CiphertextAcceptedBallot]
+    guardian: Guardian
+    context: CiphertextElectionContext
+
+
+class DecryptBallotSharesResponse(Base):
+    shares: List[BallotDecryptionShare]
+
+
 class EncryptBallotsRequest(Base):
     ballots: List[PlaintextBallot]
     seed_hash: str

app/api/v1/models/guardian.py

@@ -2,6 +2,7 @@
 
 import electionguard.election_polynomial
 import electionguard.elgamal
+import electionguard.group
 import electionguard.guardian
 import electionguard.key_ceremony
 import electionguard.schnorr
@@ -27,6 +28,7 @@
 ElectionPolynomial = Any
 ElectionPartialKeyBackup = Any
 ElectionPartialKeyChallenge = Any
+GuardianId = Any
 
 
 class Guardian(Base):
@@ -89,19 +91,25 @@ def convert_guardian(api_guardian: Guardian) -> electionguard.guardian.Guardian:
         api_guardian.number_of_guardians,
         api_guardian.quorum,
     )
+
     guardian._auxiliary_keys = electionguard.key_ceremony.AuxiliaryKeyPair(
         api_guardian.auxiliary_key_pair.public_key,
         api_guardian.auxiliary_key_pair.secret_key,
     )
+
+    election_public_key = read_json_object(
+        api_guardian.election_key_pair.public_key, electionguard.group.ElementModP
+    )
+    election_secret_key = read_json_object(
+        api_guardian.election_key_pair.secret_key, electionguard.group.ElementModQ
+    )
     guardian._election_keys = electionguard.key_ceremony.ElectionKeyPair(
+        electionguard.elgamal.ElGamalKeyPair(election_secret_key, election_public_key),
         read_json_object(
-            guardian._election_keys.key_pair, electionguard.elgamal.ElGamalKeyPair
-        ),
-        read_json_object(
-            guardian._election_keys.proof, electionguard.schnorr.SchnorrProof
+            api_guardian.election_key_pair.proof, electionguard.schnorr.SchnorrProof
         ),
         read_json_object(
-            guardian._election_keys.polynomial,
+            api_guardian.election_key_pair.polynomial,
             electionguard.election_polynomial.ElectionPolynomial,
         ),
     )