import random
import math
import sys
import struct



class RSA(object):
    """
    To generate RSA keypair we need to follow these steps:
    1. Find two large enough prime numbers, preferably of equal bytelength, p and q.
    2. Calculate modulo (p * q), it will be the common part of both private and public key.
    3. Calculate totient of p and q, also known as Euler function: phi(p, q) = (p-1)*(q-1).
    4. Find coprime number to totient and use it as public exponent, therefore completing public key.
       It can be achieved by checking elements of first 100 primes (starting from 17 for security) to have GCD = 1 with totient.
       If so, it means that these numbers are coprime. We'll use extended Euclidian algorithm.
    5. To generate private exponent, we need to find modular inverse of e modulo totient with help of EEA.
    """

    LowPrimes =   [3,5,7,11,13,17,19,23,29,31,37,41,43,47,53,59,61,67,71,73,79,83,89,97
                   ,101,103,107,109,113,127,131,137,139,149,151,157,163,167,173,179
                   ,181,191,193,197,199,211,223,227,229,233,239,241,251,257,263,269
                   ,271,277,281,283,293,307,311,313,317,331,337,347,349,353,359,367
                   ,373,379,383,389,397,401,409,419,421,431,433,439,443,449,457,461
                   ,463,467,479,487,491,499,503,509,521,523,541,547,557,563,569,571
                   ,577,587,593,599,601,607,613,617,619,631,641,643,647,653,659,661
                   ,673,677,683,691,701,709,719,727,733,739,743,751,757,761,769,773
                   ,787,797,809,811,821,823,827,829,839,853,857,859,863,877,881,883
                   ,887,907,911,919,929,937,941,947,953,967,971,977,983,991,997]


    def RabinMillerTest(self, n):
        #Miller-Rabin non-determnisctic test for primacy
        s = n - 1
        t = 0
        #Rewriting n as (2**s)*t
        while s & 1 == 0:
            s = s / 2
            t += 1
        k = 0
        while k < 128: #Selecting depth of primacy test. With k < 128, probability of n not being prime is r**(-126)
            a = random.randrange(2, n - 1)
            v = pow(a, s, n) #Very fast C hook - equal to a**s mod n
            if v != 1:
                i = 0
                while v != (n - 1):
                    if i == t - 1:
                        return False
                    else:
                        i = i + 1
                        v = pow(v, 2, n)
            k += 2
        return True

    def IsNumberPrime(self,n):
        #Checking number for primacy. To decrease computation time we'll check first 100 primes as divisors to n.
        if (n >= 3):
            if (n & 1 != 0): #We don't intrested in even numbers
                for p in self.LowPrimes:
                    if (n == p):
                        return True
                    if (n % p == 0):
                        return False
                return self.RabinMillerTest(n)
        return False

    def GeneratePrime(self,k):
        r = 100 * (math.log(k, 2) + 1)
        r_ = r
        while r > 0:
            n = random.randrange(2**(k-1),2**(k))
            r -= 1
            if self.IsNumberPrime(n) == True:
                return n
        return "Failed after "+str(r_)+" tries."

    def Totient(self,p,q):
        return (p-1)*(q-1)

    def ExtendedEuclidian(self,a,b):
        x, y,  u, v = 0, 1,  1, 0
        while a!=0:
            q, r = b//a, b%a
            m, n = x - u*q, y - v*q
            b, a, x, y, u, v = a, r, u, v, m, n
        gcd = b
        return gcd, x, y

    def ModuloInversion(self,a,m):
        gcd, x, y = self.ExtendedEuclidian(a, m)
        if gcd != 1:
            return None
        else:
            return x % m

    def GenerateKeypair(self, keylength):
        if keylength % 2 != 0 or keylength > 2048:
            print "Invalid key length (not even or too big)"
            return None
        else:
            print "OK. Generating keypair, keylength = " + str(keylength)+ " bits."
        keylength = keylength / 2
        print "Generating p and q"
        p = self.GeneratePrime(keylength)
        q = self.GeneratePrime(keylength)
        print "Calculating n..."
        n = p * q
        print "Calculating totient..."
        phi = self.Totient(p,q)
        print "Calculating e..."
        e = 0
        for prime in self.LowPrimes[5:]:
            if self.ExtendedEuclidian(prime,phi)[0]==1:
                e = prime
                break
        print "Calculating d..."
        d = self.ModuloInversion(e, phi)
        PubKey = (e, n)
        PrivKey = (d, n)
        self.WriteKey("public.txt", PubKey)
        self.WriteKey("private.txt", PrivKey)
        print "Keypair generated and saved."

    def WriteKey(self, filename, (exponent, modulo)):
        exponent = hex(exponent)[2:]
        if exponent[-1] == "L":
            exponent = exponent[:-1]

        modulo = hex(modulo)[2:]
        if modulo[-1] == "L":
            modulo = modulo[:-1]

        keytostr = exponent+":"+modulo

        with open(filename, 'w') as f:
            f.write(keytostr)

    def ReadKey(self, filename):
        with open(filename, 'r') as f:
            strtokey = f.read().split(':')
            exponent = int(strtokey[0],16)
            modulo = int(strtokey[1],16)
            return exponent, modulo


    def Encode(self, pubkeypath, filepath):
        exponent, modulo = self.ReadKey(pubkeypath)
        keylen = (len(bin(modulo))-2)/8
        chunksize = keylen - 1

        with open(filepath, 'rb') as f:
            filecontent = f.read()

        #Break file in chunks sized keylen/8 - 1 to encode'
        chunks , lastsize = self.Chunker(filecontent, chunksize)

        #Encode:
        processedchunks = [pow(chunks[x],exponent, modulo) for x in xrange(len(chunks))]
        processedcontent = self.DechunkerEnc(processedchunks, chunksize)

        header = "SNKKZRSA"+str(lastsize)+'&&&&&&&&&&'
        with open(filepath, 'wb') as f:
            f.write(header+processedcontent)


    def Decode(self, privkeypath, filepath):
        exponent, modulo = self.ReadKey(privkeypath)
        keylen = (len(bin(modulo))-2)/8
        chunksize = keylen

        with open(filepath, 'rb') as f:
            filecontent = f.read()

        #Extract header
        headerstop=filecontent.find("&&&&&&&&&&")
        header = filecontent[:headerstop+10]
        lastsize = int(header[8:-10])
        filecontent=filecontent[headerstop+10:]

        #Break file in chunks sized keylen/8  to decode'
        chunks = self.Chunker(filecontent, chunksize)[0]

        #Decode:
        processedchunks = [pow(chunks[x],exponent, modulo) for x in xrange(len(chunks))]
        processedcontent = self.DechunkerDec(processedchunks, chunksize, lastsize)

        with open(filepath, 'wb') as f:
            f.write(processedcontent)



    def Chunker(self, st, chunksize):
        chunks = [st[x:x+chunksize] for x in xrange(0,len(st),chunksize)]
        intchunks = [int(chunks[x].encode('hex'), 16) for x in xrange(len(chunks))]
        return intchunks, len(chunks[-1])

    def DechunkerEnc(self, chunks, chunksize):
        chunksize += 1
        hexchunks = [hex(chunks[x])[2:] for x in xrange(len(chunks))]

        for x in xrange(len(hexchunks)):
            if hexchunks[x][-1] == "L":
                hexchunks[x] = hexchunks[x][:-1]
            if len(hexchunks[x])<chunksize*2:
                hexchunks[x] = '0'*((chunksize*2)-len(hexchunks[x])) + hexchunks[x]
        dechunked = ''
        for x in xrange(len(hexchunks)):
            dechunked+=hexchunks[x].decode('hex')

        return dechunked

    def DechunkerDec(self, chunks, chunksize, lastsize):
        chunksize -= 1
        hexchunks = [hex(chunks[x])[2:] for x in xrange(len(chunks))]

        for x in xrange(len(hexchunks)):
            if hexchunks[x][-1] == "L":
                hexchunks[x] = hexchunks[x][:-1]

        for x in xrange(len(hexchunks)-1):
            if len(hexchunks[x])<chunksize*2:
                hexchunks[x] = '0'*((chunksize*2)-len(hexchunks[x])) + hexchunks[x]

        if len(hexchunks[-1]) < lastsize*2:
            hexchunks[-1]='0'*((lastsize*2)-len(hexchunks[-1]))+hexchunks[-1]


        dechunked = ''
        for x in xrange(len(hexchunks)):
            dechunked+=hexchunks[x].decode('hex')

        return dechunked




usage = "RSA cipherer and key generator. \nUsage:\n     RSA2.py generate <keylength> - generate key pair.\n     RSA2.py encrypt <public key file> <file to encrypt> - Encrypt a file.\n     RSA2.py decrypt <private key file> <file to decrypt> - Decrypt a file."

asymmetrical = RSA()

if len(sys.argv) < 2 or len(sys.argv) > 4 or sys.argv[1] == 'help':
    print usage

try:
    if sys.argv[1] == 'generate':
        keylen=int(sys.argv[2])
        asymmetrical.GenerateKeypair(keylen)
        quit()
    elif sys.argv[1] == 'encrypt':
        keypath=sys.argv[2]
        filepath=sys.argv[3]
        asymmetrical.Encode(keypath,filepath)
        quit()
    elif sys.argv[1] == 'decrypt':
        keypath=sys.argv[2]
        filepath=sys.argv[3]
        asymmetrical.Decode(keypath,filepath)
        quit()
except IndexError:
    print "Invalid parameters"