🐞 Shared selection IDs across contests silently corrupt the batch tally #816

import logging
from datetime import datetime, timezone

from electionguard.ballot import BallotBoxState, PlaintextBallot, PlaintextBallotContest, PlaintextBallotSelection
from electionguard.ballot_box import BallotBox, get_ballots
from electionguard.data_store import DataStore
from electionguard.decryption_mediator import DecryptionMediator
from electionguard.election import make_ciphertext_election_context
from electionguard.encrypt import encrypt_ballot
from electionguard.group import ONE_MOD_Q
from electionguard.guardian import Guardian
from electionguard.key_ceremony import CeremonyDetails
from electionguard.key_ceremony_mediator import KeyCeremonyMediator
from electionguard.manifest import (
    BallotStyle, Candidate, ContestDescription, ElectionType, GeopoliticalUnit, InternalManifest,
    Manifest, ReportingUnitType, SelectionDescription, SpecVersion, VoteVariationType,
)
from electionguard.tally import tally_ballots

logging.disable(logging.INFO)  # keep the library's warnings, drop its INFO lines


def selection(object_id, sequence_order):
    return SelectionDescription(object_id, sequence_order, candidate_id=object_id)


def contest(object_id, sequence_order, selections, number_elected=1, votes_allowed=1):
    return ContestDescription(object_id, sequence_order, "district", VoteVariationType.one_of_m,
                              number_elected, votes_allowed, object_id, selections)


def make_manifest(contests, styles=None, units=("district",)):
    candidates = sorted({s.candidate_id for c in contests for s in c.ballot_selections})
    return Manifest(
        "scope", SpecVersion.EG1_0, ElectionType.general,
        datetime(2020, 3, 1, tzinfo=timezone.utc), datetime(2020, 3, 2, tzinfo=timezone.utc),
        [GeopoliticalUnit(u, u, ReportingUnitType.county) for u in units], [],
        [Candidate(c) for c in candidates], contests,
        styles or [BallotStyle("style", ["district"])],
    )


def ballot(object_id, marks, style_id="style"):
    """marks: {contest_id: [selected selection ids]}"""
    return PlaintextBallot(object_id, style_id, [
        PlaintextBallotContest(c, [PlaintextBallotSelection(s, 1) for s in sel]) for c, sel in marks.items()])


def run_election(manifest, ballots, spoil=()):
    """One guardian, quorum 1: key ceremony, encrypt, cast or spoil, tally, decrypt."""
    print("manifest.is_valid():", manifest.is_valid())
    guardian = Guardian.from_nonce("guardian-1", 1, 1, 1)
    key_mediator = KeyCeremonyMediator("key-mediator", CeremonyDetails(1, 1))
    key_mediator.announce(guardian.share_key())
    joint_key = key_mediator.publish_joint_key()
    internal = InternalManifest(manifest)
    context = make_ciphertext_election_context(
        1, 1, joint_key.joint_public_key, joint_key.commitment_hash, manifest.crypto_hash())

    store = DataStore()
    box = BallotBox(internal, context, store)
    for plain in ballots:
        encrypted = encrypt_ballot(plain, internal, context, ONE_MOD_Q, should_verify_proofs=True)
        if encrypted is None:
            print(f"{plain.object_id}: encrypt_ballot returned None")
            continue
        submitted = box.spoil(encrypted) if plain.object_id in spoil else box.cast(encrypted)
        print(f"{plain.object_id}: {'spoil' if plain.object_id in spoil else 'cast'} ->",
              None if submitted is None else submitted.state.name)

    tally = tally_ballots(store, internal, context)
    spoiled = list(get_ballots(store, BallotBoxState.SPOILED).values())
    decryption = DecryptionMediator("decryption-mediator", context)
    decryption.announce(guardian.share_key(), guardian.compute_tally_share(tally, context),
                        guardian.compute_ballot_shares(spoiled, context))
    plaintext = decryption.get_plaintext_tally(tally, manifest)
    for c in plaintext.contests.values():
        print("tally", c.object_id, {s.object_id: s.tally for s in c.selections.values()})
    if spoil:
        decrypted = decryption.get_plaintext_ballots(spoiled, manifest)
        print("decrypted spoiled ballots:", sorted(decrypted))
    return plaintext

# Two ballot measures; each has selections with object_id "yes" and "no".
manifest = make_manifest([
    contest("measure-1", 0, [selection("yes", 0), selection("no", 1)]),
    contest("measure-2", 1, [selection("yes", 0), selection("no", 1)]),
])
run_election(manifest, [
    ballot("b1", {"measure-1": ["yes"], "measure-2": ["no"]}),
    ballot("b2", {"measure-1": ["yes"], "measure-2": ["no"]}),
    ballot("b3", {"measure-1": ["no"], "measure-2": ["no"]}),   # spoiled
], spoil={"b3"})
# Expected: manifest.is_valid() is False, since "yes" and "no" appear in both contests.
# Instead it is True, and measure-1 should be {yes: 2, no: 0} but is reported as {yes: 0, no: 2}.