""" Class defintion for common factor exploit """

sage_present = True
try:
    from sage.all_cmdline import *
except ImportError:
    sage_present = False

from RSAExploits.util import compute_priv_rsa_exp
from fractions import gcd
from RSAExploits.exploits.exploit import Exploit

# Used to hopefully limit calls to sage's Primes() function
primes_seen = []

class Common_Factor(Exploit):
    """ Class providing a run interface to the common factor exploit"""

    def run(self, rsadata_list, info_dict = None):
        """ Attempts to recover plaintext using the common factor exploit

            This attack works when we have 2 RSA_Data objects with moduli
            that share a common factor and one of the moduli is the product
            of two primes
            
            Args:
                rsadata_list: A list of RSA_Data objects with moduli to search for
                              common factors
                info_dict: Not currently used

            Side Effects: 
                If a modulus is factored, the corresponding RSA_Data object's
                private key will be computed and initialized

            Return: 
                True if at least one modulus is factored
        """
        
        global sage_present

        if not sage_present:
            print "Common Factor: Sage required, skipping exploit"
            return 
        
        print "Common Factor: Running Attack..."
        success = False
        for i in range(0, len(rsadata_list) - 1):
            for j in range(1, len(rsadata_list)):
                if self.common_factor_exploit(rsadata_list[i], rsadata_list[j]):
                    success = True

        if success:
            print "Common Factor: Success, at least one private exponent recovered"
        else:   
            print "Common Factor: Failure, no private exponents recovered"
        return success

    @staticmethod
    def common_factor_exploit(rsadata1, rsadata2):
        """ Attempt to recover the private keys of the RSA_Data objects

            Checks if the moduli in the two rsadata objects share a common factor.
            If so, this function will attempt to recover the private exponent
            used for decryption. Returns true if at least one private exponent 
            has been recovered. Updates the d, p, and q fields of an rsadata
            object if the exploit is successful
        """
    
        global primes_seen

        ret =False

        # If the two moduli have no common factors
        p = long(gcd(rsadata1.get_n(), rsadata2.get_n()))
        if p == 1:
            return ret

        # Compute a second factor for each modulus
        q1 = long(rsadata1.get_n() / p)
        q2 = long(rsadata2.get_n() / p)
        
        # If the two factors of the first modulus are both prime
        if (p in primes_seen or p in Primes()) and (q1 in primes_seen or q1 in Primes()):

            # Update primes list to hopefully save time in the future
            primes_seen.append(p)
            primes_seen.append(q1)

            # Recover d and update fields in rsadata object
            d = compute_priv_rsa_exp(rsadata1.get_e(), (p, q1))
            rsadata1.set_d(d)
            ret = True

        # If the two factors of the second modulus are both prime
        if (p in primes_seen or p in Primes()) and (q2 in primes_seen or q2 in Primes()):
            primes_seen.append(p)
            primes_seen.append(q2)
            d = compute_priv_rsa_exp(rsadata2.get_e(), (p, q2))
            rsadata2.set_d(d)
            ret = True

        return ret

