# pre_election_audit.py
# Ronald L. Rivest
# October 11, 2009
# 
# This Python program is for use with the Scantegrity II election system.
# See www.scantegrity.org for information on Scantegrity II.
# 
# It was written specifically for the version of Scantegrity II used
# in the Takoma Park, Maryland election of November 3rd, 2009.  It
# may need slightly modifications for use with other versions of
# Scantegrity.
#
# This program is invoked as:
#    python  pre_election_audit.py  randomDataFileName  inFileName  outFileName percentage_challenged
#       randDataFileName must be given; it is arbitrary file name (e.g. stock price data)
#       if inFileName is omitted, it is MeetingOneIn.xml
#       if outFileName is omitted, it is MeetingTwoIn.xml
#       percentage_challenged should be an integer between 0 and 100, inclusive
#       if percentage_challenged is omitted, it is 50 (i.e. 50% of ballots will be audited)
#
#   The randomDataFile may contain arbitrary data; in practice it
#   is expected to contain closing stock price data, as produced by 
#       get_latest_djia_stock_prices.py
#   which produces the file
#       djia_stock_prices_latest.txt
#    
#   This program:
#
#    (1) Reads an XML parameter specification file named
#            MeetingOneIn.xml      (or something else if specified on command line)
#        which will look something like:
#
#        <xml>
#          <electionSpec>D:\PunchScan2.0\PunchScan2.0\Elections\TP Nov 3 2009, mock\ward2\ElectionSpec.xml</electionSpec>
#          <noBallots>100</noBallots>
#          <noDs>1</noDs>
#          <constant>d2FyZDIgICAgICAgICAgIA==</constant>
#          <partitions>D:\PunchScan2.0\PunchScan2.0\Elections\TP Nov 3 2009, mock\ward2\partitions.xml</partitions>
#        </xml>
#
#    (2) Obtains the following parameters from that file:
#            noBallots         (nonnegative integer)
#            constant          (base-64 encoded election constant)
#
#    (3) Produces an output file called
#            MeetingTwoIn.xml       (or something else if specified on command line)
#        containing
#
#        <xml>
#          <challenges>
#            <print>
#               ROWDATA
#            </print>
#          </challenges>
#        </xml>
#
#        where ROWDATA is a number of distinct rows of the form: 
#            <row id="84"/>
#        where each row id is between 0 and noBallots-1
#        The number of rows is equal to the specified number of challenges.

#
# IMPORTS
#
import base64
import hashlib
import sys
from xml.etree import ElementTree as etree

#
# ROUTINE FOR READING XML INPUT FILE
#

def read_MeetingOneIn(inFileName):
    """
    Read the XML input file inFileName, and
    return the values of the  noBallots  and  constant  found therein.
    """
    doc = etree.parse(inFileName)
    root = doc.getroot()
    for node in root:
        if node.tag == "noBallots":
            noBallots = int(node.text)
        if node.tag == "constant":
            constant = base64.decodestring(node.text);
    return (noBallots,constant)

#
# ROUTINE FOR PRODUCING OUTPUT FILE
#

def write_MeetingTwoIn(outFileName,challenge_list):
    file = open(outFileName,"w")
    file.write("<xml><challenges><print>\n")
    for row in challenge_list:
        file.write('<row id="%d"/>\n'%row)
    file.write("</print></challenges></xml>\n")
    file.close()

def prn(seed,index,modulus):
    """ 
    Return pseudo-random number based on seed, index, and modulus:
         seed is arbitrary byte string
         index is integer to distinguish different calls
         modulus specifies what the output modulus is; value is between 0 and modulus-1
    """

    # Append decimal representation of index at end of seed to form hash input
    hash_input = seed + "%d"%index
    
    # Apply SHA1 hash function to get a byte sequence of length 20 bytes (160 bits),
    # encoded as 40 hex digits.
    hash_output = hashlib.sha1(hash_input).hexdigest()

    # Return hash output (treated as integer) modulo given modulus
    hash_value = int(hash_output,16)      # interpret hash output as a base-16 integer
    output_value = hash_value % modulus   # mod modulus ==> output value
    return output_value

def main():
    args = sys.argv[1:]
    print
    print "    Scantegrity II Pre-Election Audit Challenge Generator"
    print "    Version 1.0. October 11, 2009."
    print 

    if len(args) == 0:
        print '     usage: python pre_election_audit.py randomDataFileName inFileName outFileName percentage_challenged'
        print "            randDataFileName must be given; it is arbitrary file name (e.g. stock price data)"
        print "            if inFileName is omitted, it is MeetingOneIn.xml"
        print "            if outFileName is omitted, it is MeetingTwoIn.xml"
        print "            percentage_challenged should be between 0 and 100, inclusive"
        print "            if percentage_challenged is omitted, it is 50 (i.e. 50% of ballots will be audited)"
        print 
        sys.exit(-1)

    randomDataFileName = args[0]

    if len(args) >= 2:
        inFileName = args[1]
    else:
        inFileName = "MeetingOneIn.xml"

    if len(args) >= 3:
        outFileName = args[2]
    else:
        outFileName = "MeetingTwoIn.xml"

    if len(args) >= 4:
        percentage_challenged = int(args[3])
        percentage_challenged = min(100,max(0,percentage_challenged))
    else:
        percentage_challenged = 50

    # read random data seed from file
    print "    Step 1. Obtain random data seed from:",randomDataFileName
    f = open(randomDataFileName,"rb")
    seed = f.read()
    f.close()

    # Read XML input (MeetingOneIn.xml file by default)
    (noBallots,constant) = read_MeetingOneIn(inFileName)
    print "    Step 2. Read input file:", inFileName
    print "            Number of ballots = ",noBallots
    print "            constant = '",constant,"'"

    # set noChallenges
    noChallenges = (percentage_challenged * noBallots) // 100
    noChallenges = min(noBallots,noChallenges)
    print "    Step 3. Determine number of challenges desired = ",noChallenges

    # append election constant to seed
    seed = seed + constant

    # make up list of challenges
    print "    Step 4. Generate challenges."
    challenge_list = []
    challenge_array = [0]*noBallots                  # to mark ones already taken
    counter = 0                                      # counts number of prn calls
    for i in range(noChallenges):                    # until we have desired number challenges
        while (challenge_array[prn(seed,counter,noBallots)]==1):
            counter += 1
        index = prn(seed,counter,noBallots)          # mark this row index as now used
        challenge_array[index] = 1
        challenge_list.append(index)                 # and save it on list of challenges
    challenge_list = sorted(challenge_list)          # sort into order (prettier output)
    print "            Number of challenges produced =", len(challenge_list)

    # save output in desired format
    print "    Step 5. Write challenge list to output file:",outFileName
    print
    write_MeetingTwoIn(outFileName,challenge_list)

    print "    Done."
    print

if __name__ == '__main__':
    main()