import hashlib

def disjprove1(p,q,R,S,g,pk,r): #taken from exercise 2, proof when m=1
    ch0 = ZZ.random_element(q);
    rep0 = ZZ.random_element(q);
    U0 = Mod(power_mod(g,rep0,p) * inverse_mod(power_mod(R,ch0,p),p),p);
    V0 = Mod(power_mod(pk,rep0,p) * inverse_mod(power_mod(S,ch0,p),p),p);
    u_1 = ZZ.random_element(q);
    U1 = power_mod(g,u_1,p);
    V1 = power_mod(pk,u_1,p);
    hash_value = hashlib.sha256(str(R) + str(S) + str(U0) + str(V0) + str(U1) + str(V1)).hexdigest();
    ch = Integer(int(hash_value,16),q);
    ch1 = Integer(Mod((ch - ch0),p));
    rep1 = Integer(Mod(u_1 + r * ch1,q));
    return [ch0,ch1,rep0,rep1];

def disjprove0(p,q,R,S,g,pk,r): #taken from exercise 2, proof when m=0
    u_0 = ZZ.random_element(q);
    U0 = power_mod(g,u_0,p);
    V0 = power_mod(pk,u_0,p);
    ch1 = ZZ.random_element(q);
    rep1 = ZZ.random_element(q); 
    U1 = Mod(power_mod(g,rep1,p) * inverse_mod(power_mod(R,ch1,p),p),p);
    V1 = Mod(power_mod(pk,rep1,p) * inverse_mod(power_mod((S/g),ch1,p),p),p);
    hash_value = hashlib.sha256(str(R) + str(S) + str(U0) + str(V0) + str(U1) + str(V1)).hexdigest();
    ch = Integer(int(hash_value,16),q);
    ch0 = Integer(Mod(ch-ch1,p));
    rep0 = Integer(Mod(u_0 + r * ch0,q));
    return [ch0,ch1,rep0,rep1];

def disjverify(p,q,g,ch0,ch1,rep0,rep1,R,S): #taken from exercise 2, verify
    val1 = Mod(power_mod(g,rep0,p) * inverse_mod(power_mod(R,ch0,p),p),p);
    val2 = Mod(power_mod(pk,rep0,p) * inverse_mod(power_mod(S,ch0,p),p),p);      
    val3 = Mod(power_mod(g,rep1,p) * inverse_mod(power_mod(R,ch1,p),p),p); 
    val4 = Mod(power_mod(pk,rep1,p) * inverse_mod(power_mod((S*power_mod(g,-1,p)),ch1,p),p),p); 
    hash_value = hashlib.sha256(str(R) + str(S) + str(val1) + str(val2) + str(val3) + str(val4)).hexdigest();
    val5 = Mod(ch0+ch1,p);

    return val5 == Integer(int(hash_value,16),q)

def eqdlprove(p,q,g1,g2,y1,y2,x):
    r = ZZ.random_element(q);
    com1 = power_mod(g1,r,p);
    com2 = power_mod(g2,r,p);
    hash_value = hashlib.sha256(str(y1) + str(y2) + str(com1) + str(com2)).hexdigest();
    ch = Integer(int(hash_value,16),q);
    s = Integer(Mod(r + x * ch,q));
    return [ch,s]

def setup(): #setup algorithm
    #Assume a value for p
    p = 19557005198709002903944860293031210393112763219086296641657077101075696217230957575932867920574816565856404011416908876733063974554042055901449481509845537165833627083609136426386625960849573203936690977282752217830467255076903813033768202255463477317279507253420233052258146393806024179130344501759742626253484354878829682461413002143954786668386100202257237258007726886368228024987052808444032800050116408995869561195208576647050961019393809940446490709061546503439488751331252440431940049716566765873615456614911458658341732120931319694200689362617296630728389003180670326151653001322354906363854478585334821235883
    #Get value of q
    q = Integer((p-1)/2);
    #Assume a value for g
    g = 4;
    #Generate secret key by choosing a random value in Zq
    sk = ZZ.random_element(q);
    #Get public key by calculating g^sk
    pk = power_mod(g,sk,p);
    return[p,q,g,sk,pk];

def vote(voter_id, v): #vote algorithm
    r = ZZ.random_element(q);
    R = power_mod(g, r, p); #compute R
    S = Integer(Mod(power_mod(pk,r,p)*power_mod(g,v,p),p)); #compute S
    if(v==1):
        [ch0,ch1,rep0,rep1] = disjprove1(p,q,R,S,g,pk,r);
    else:
        [ch0,ch1,rep0,rep1] = disjprove0(p,q,R,S,g,pk,r);
    #output R,S and disjunction prover values (pi on notes)
    return [R,S,ch0,ch1,rep0,rep1];

def validate(R,S,ch0,ch1,rep0,rep1): #validate algorithm - needed for tally algorithm
    if disjverify(p,q,g,ch0,ch1,rep0,rep1,R,S) == True:
        return True;
    else:
        return False;

def tally(BB,sk):
    R_tot = 1; #R and S need to be 1 since we are multiplying each time. If 0, final result would be 0 (0*x =0)
    S_tot = 1;
    g_result = 0;
    result = 0; #final result of the election
    for b in BB: #loop over every ballot in the box - check if the vote is valid
        R = b[0]; S = b[1]; ch0 = b[2]; ch1 = b[3]; rep0 = b[4]; rep1 = b[5];
        if validate(R,S,ch0,ch1,rep0,rep1) == False: #invalid vote
            return "invalid, pi := 0";
        else: #valid vote - multiply R_tot and S_tot to get total votes - homomorphic property
            R_tot = Mod(R_tot * R,p);
            S_tot = Mod(S_tot * S,p);
    g_result = Mod(S_tot/(power_mod(R_tot,sk,p)),p); #compute g^result
    [ch,s] = eqdlprove(p,q,g,R_tot,pk,Mod(S_tot/g_result,p), sk); #ZK proof of Equality of Logarithms
    i=0;

    #get result by using pre-computed array
    for r in result_arr:
        if g_result == r:
            result = i;
        i += 1;
    return [ch,s,result];

#Building a toy referendum election with 10 voters

no_voters = 10; #given number of voters
[p,q,g,sk,pk] = setup(); #setup
BB=[];
#pre-computed results - we can do this because we know the number of voters
result_arr = [power_mod(g,0,p),power_mod(g,1,p),power_mod(g,2,p),power_mod(g,3,p),power_mod(g,4,p),power_mod(g,5,p),power_mod(g,6,p),power_mod(g,7,p),power_mod(g,8,p),power_mod(g,9,p),power_mod(g,10,p)];

for voter_id in range(1,11):
    v = ZZ.random_element(0,2); #get random vote/ballot - 1 or 0
    print("voter_id: "+str(voter_id)+", vote: "+str(v));
    b=vote(voter_id, v); #cast vote
    if validate(b[0],b[1],b[2],b[3],b[4],b[5]):
        BB.append(b); #add to Ballot Box, if validate(b) accepts

print("\n");
[ch,s,result] = tally(BB,sk);
print("challenge = " + str(ch) + "\n")
print("solution = " + str(s) + "\n")
print("result = " + str(result) + "; Therefore "+str(result)+" people voted 1 and "+str(no_voters-result)+" people voted 0.\n")