""" The class defintion for RSAData """
    
from Crypto.PublicKey import RSA
import fractions 
import random

class RSAData:
    """ Conttains all RSA information used by currently implemented exploits

        RSAData is built on top of two other classes. The first is the RSAObj
        taken from the pycrypto library. This class stores the RSA modulus n,
        public exponent e, private exponent d, and prime factors p and q. We
        currently leverage the decrypt method from RSAObj. The second class is
        TextData which currently holds the ciphertext/plaintext
        information.

        Unfortunately the underlying implementation of RSAData is not
        completely hidden from the user given the way the __init__ function
        receives parameters. However, from a user's perspective, RSAData
        should be a container that has the following fields (along with an
        encrypt and decrypt function):

        n: RSA modulus
        e: Public exponent
        d: Private exponent
        p: First prime factor
        q: Second prime factor
        c: Ciphertext
        m: Plaintext
        id_num: Used to specify which ciphertexts are generated from the
                same plaintext
    """

    def __init__(self, rsaobj, txt=None):
        """ Initialize the 2 classes using the provided parameters """
        self.rsaobj = rsaobj
        self.txt = txt
	
    def encrypt(self):
        """ Encrypt the current plaintext and assign it to the ciphertext """
        if self.get_m() != None:
            c_text = self.rsaobj.encrypt(long(self.get_m()))
            self.set_c(c_text)	
            return True
        else:
            return False
            
    def decrypt(self):
        """ Encrypt the current ciphertext and assign it to the plaintext """
        if self.get_c() != None and self.rsaobj.has_private():
            m_text = self.rsaobj.decrypt(self.get_c())
            self.set_m(m_text)	
            return True 
        else:
            return False

    def get_e(self):
        """ Return the public exponent """
        if hasattr(self.rsaobj, "e"):
            return self.rsaobj.e
        else:
            return None

    def set_e(self, key_e):
        """ Set the public exponent """
        self.rsaobj = rsa_obj(self.get_n(), key_e, self.get_d(), \
                            self.get_p(), self.get_q())

    def get_p(self):
        """ Return the first prime factor """
        if hasattr(self.rsaobj, "p"):
            return self.rsaobj.p
        else:
            return None
    
    def set_p(self, key_p):
        """ Set the first prime factor """
        self.rsaobj = rsa_obj(self.get_n(), self.get_e(), \
                              self.get_d(), key_p, self.get_q())

    def get_q(self):
        """ Return the second prime factor """
        if hasattr(self.rsaobj, "q"):
            return self.rsaobj.q
        else:
            return None
    
    def set_q(self, key_q):
        """ Set the second prime factor """
        self.rsaobj = rsa_obj(self.get_n(), self.get_e(), \
                              self.get_d(), self.get_p(), key_q)

    def get_n(self):
        """ Return the modulus """
        if hasattr(self.rsaobj, "n"):
            return self.rsaobj.n
        else:
            return None
    
    def set_n(self, key_n):
        """ Set the modulus """
        self.rsaobj = rsa_obj(key_n, self.get_e(), self.get_d(), \
                              self.get_p(), self.get_q())
    
    def get_d(self):
        """ Return the private exponent """
        if hasattr(self.rsaobj, "d"):
            return self.rsaobj.d
        else:
            return None

    def set_d(self, key_d):
        """ Set the private exponent """
        self.rsaobj = rsa_obj(self.get_n(), self.get_e(), key_d, \
                              self.get_p(), self.get_q())

    def get_c(self):
        """ Return the ciphertext """
        if self.txt == None:
            return None
        return self.txt.c
    
    def set_c(self, c_text):
        """ Set the ciphertext """
        if self.txt == None:
            self.txt = TextData(c=c_text)
            return

        self.txt = TextData(c_text=c_text, idnum=self.get_id(), \
                            m_text=self.get_m())

    def get_m(self):
        """ Return the plaintext """
        if self.txt == None:
            return None
        return self.txt.m
    
    def set_m(self, m_text):
        """ Set the plaintext """
        if self.txt == None:
            self.txt = TextData(m_text=m_text)
            return
        self.txt = TextData(c_text=self.get_c(), idnum=self.get_id(), \
                            m_text=m_text)

    def get_id(self):
        """ Return the id """
        if self.txt == None:
            return None
        return self.txt.idnum
    
    def set_id(self, idnum):
        """ Set the id """
        if self.txt == None:
            self.txt = TextData(idnum=idnum)
            return
        self.txt = TextData(c_text=self.get_c(), idnum=idnum, \
                            m_text=self.get_m())


def rsa_obj(key_n, key_e, key_d=None, key_p=None, key_q=None):
    """ Wrapper for the RSAObj constructor
	
        The main reason for its existance is to compute the
        prime factors if the private exponent d is being set.
        In testing, the construct method threw exceptions because
        it wasn't able to compute the prime factors. The 
        recover_prime_factors function seems work better.
    """	
    if key_n != None and key_e != None and key_d == None \
    and key_p == None and key_q == None:
        key = RSA.construct((key_n, key_e))
    elif key_n != None and key_e != None and key_d != None \
    and key_p == None and key_q == None:
        key_p, key_q = recover_prime_factors(key_n, key_e, key_d)		
        key = RSA.construct((key_n, key_e, key_d, long(key_p), long(key_q)))
    elif key_n != None and key_e != None and key_d != None \
    and key_p != None and key_q == None:				
        key = RSA.construct((key_n, key_e, key_d, key_p, key_n/key_p))	
    elif key_n != None and key_e != None and key_d != None \
    and key_p != None and key_q != None:
        key = RSA.construct((key_n, key_e, key_d, key_p, key_q))
    return key	

class TextData:
    """ At the moment, a container for the cipher/plaintext information """
    def __init__(self, c_text=None, idnum=None, m_text=None):
        """ Standard initializer """
        self.c = c_text
        self.idnum = idnum
        self.m = m_text

 
def fail_function():
    """ Basic fail function for recover_prime_factors() """
    print("Prime factors not found")
 
def output_primes(a_num, n_num):
    """ Returns the recovered prime factors """
    p_num = fractions.gcd(a_num, n_num)
    q_num = int(n_num / p_num)
    if p_num > q_num:
        p_num, q_num = q_num, p_num
    return p_num, q_num
 
 
def recover_prime_factors(key_n, key_e, key_d):
    """The following algorithm recovers the prime factor
        s of a modulus, given the public and private
        exponents.
        Function call: RecoverPrimeFactors(n, e, d)
        Input:  n: modulus
                e: public exponent
                d: private exponent
        Output: (p, q): prime factors of modulus"""
 
    k_num = key_d * key_e - 1
    if k_num % 2 == 1:
        fail_function()
        return 0, 0
    else:
        t_num = 0
        r_num = k_num
        while(r_num % 2 == 0):
            r_num = int(r_num / 2)
            t_num += 1
        for i_num in range(1, 101):
            g_num = random.randint(0, key_n) # random g in [0, n-1]
            y_num = pow(g_num, r_num, key_n)
            if y_num == 1 or y_num == key_n - 1:
                continue
            else:
                for j_num in range(1, t_num): # j \in [1, t-1]
                    x_num = pow(y_num, 2, key_n)
                    if x_num == 1:
                        key_p, key_q = output_primes(y_num - 1, key_n)
                        return key_p, key_q
                    elif x_num == key_n - 1:
                        continue
                    y_num = x_num
                    x_num = pow(y_num, 2, key_n)
                    if  x_num == 1:
                        key_p, key_q = output_primes(y_num - 1, key_n)
                        return key_p, key_q

