#!/usr/bin/env python

import pbp #Get access to encrypt and decrypt functions
from Crypto.Util import number #For modular multiplicative inverse
import math
from Crypto.PublicKey import RSA
#Had to use "sudo apt-get install python-dev" to be able to upgrade pycrypto with "sudo pip install pycrypto --upgrade"

MODULI_FILE = 'moduli.hex' #'test.hex' #
CIPHER_FILE = '1.2.4_ciphertext.enc.asc'
PRIMES_FILE = 'primes.txt'
OUTPUT_FILE = 'sol_1.2.4.txt'
e = 65537 #Public exponent

#In python long types should have infinite precision besides memory limits
#Get RSA moduli (N = p * q, p and q are prime factors)
with open(MODULI_FILE,'r') as mod_f:
    RSA_mods = mod_f.read()
    RSA_mods = RSA_mods.strip()
    RSA_mods = RSA_mods.split('\n')

RSA_ints = []
#Convert array of hex moduli to ints (might be "long")
for i in range(len(RSA_mods)):
    RSA_ints.append(int(RSA_mods[i],16))

#Get cipher text
with open(CIPHER_FILE, 'r') as ciph_f:
    ciph = ciph_f.read()


#Compute pairwise GCDs of RSA keys
#Multiplies the first two entries in a vector
def prod(a):
    p = 1
    for i in range(len(a)):
        p = p * a[i]
    return p


    #Product tree: using method from https//:facthacks.cr.yp.to/product.html
    #Takes in vector of integers to multiply together (will be the RSA moduli)
def producttree(X):
    result= [X]
    while len(X) > 1:
        X = [prod(X[i*2:(i+1)*2]) for i in range ((len(X)+1)/2)]
        result.append(X)
    return result

def product(X):
    if len(X) == 0: return 1
    while len(X) > 1:
        X = [prod(X[i*2:(i+1)*2]) for i in range((len(X) + 1)/2)]
    return X[0]

    #Remainder tree: using method from https//:facthacks.cr.yp.to/remainder.html
    #def remainderusingproducttree(n,T):
    #    result = [n]
    #    for t in reversed(T):
    #        #Floor function on i/2 implied for ints
    #       #print T[len(T)-10: len(T)-1] #Errors?
    #        result = [result[(i/2)]%t[i] for i in range(len(t))]

    #    return result

    #Computes n mod X[i] for i in range(len(X))
    #def remainders(n,X):
    #    return remainderusingproducttree(n,producttree(X))

    #Not defined for python 2 it seems
def gcd(a, b):
        #Check that a > b
    if(a < b):
        c = a
        a = b
        b = c
    while( b > 0):
        c = a%b
        a = b
        b = c
    return a

def batchgcd_faster(X):
    prods = producttree(X)
    R = prods.pop()
    while prods:
        X = prods.pop()
        R = [R[int(math.floor(i/2))] % X[i]**2 for i in range(len(X))]
    return [gcd(r/n,n) for r,n in zip(R,X)]

    #Actual call over RSA moduli, need to take gcd(product of keys mod (Ni^2)/Ni,Ni) for
    #all keys Ni

    #final_prod = product(RSA_ints)
    #RSA_squared = []
    #for i in range(len(RSA_ints)):
    #    RSA_squared.append(RSA_ints[i]**2)
    #RSA_rem = remainders(final_prod, RSA_squared)
    #del RSA_squared
pq = []
RSA_gcd = batchgcd_faster(RSA_ints)

for i in range(len(RSA_ints)):
        #Find gcd of results from remainder tree and product tree
        #RSA_rem[i] = gcd(RSA_rem[i]/RSA_ints[i], RSA_ints[i])
        #Looking through remainder tree results for gcds that resulted in non-1 values (no
        #factors matched) as well as values that that don't match the RSA key (both factors
        #matched) and appending to a final list of keys and one of their factors
    if(RSA_gcd[i] != 1 and RSA_gcd[i] != RSA_ints[i]):
            #Factor RSA keys with common factors by dividing and save the two factors for
            #easily finding Totient function to decrypt
            #Broken keys will be of form [p,q,i] where p is found shared factor, q is the
            # other factor found bu dividing the key by the found factor p, and i is the
            #number of the broken key in the original
        pq.append([RSA_gcd[i], RSA_ints[i] / RSA_gcd[i], i])

with open(PRIMES_FILE,'w') as prime_f:
    for i in range(len(pq)):
        prime_f.write(str(pq[i]) + '\n')

    #Decrypt cipher using RSA keys
    #Find private key, not sure if I need to include k, not for now
    # m^(k*fi(n) + 1)  = m (mod n) whfiere e*d = k*fi(n) + 1, d = (k*fi(n) + 1)/e
    #(q-1)*(p-1) = fi(N) = fi(q)*(p)
#Assume k = 1
#number.inverse(e,p*q-1) finds d such that e*(p*q-1)  = 1 (mod b)
with open(OUTPUT_FILE,'w') as output:
#output = open("sol_1.2.4.txt", 'w')
    d = [] #Broken keys
    for i in range(len(pq)):
        #d = ((broken_keys[i][0]-1)*(broken_keys[i][1]-1) + 1)/e
        d.append(number.inverse(e, (pq[i][0]-1)*(pq[i][1]-1)))
        #Output plaintext to file, outputting all possibilities and will check for
        #For clue manually, then output the correct plaintext to solution
        tup = (long(pq[i][0]*pq[i][1]), long(e), long(d[i]))
        RSA_key = RSA.construct(tup)
        try:
            plain =  pbp.decrypt(RSA_key,cipher)
            print plain
            output.write(plain)
        except ValueError:
            pass
#output.close()