""" Class defintion for fermat factorization """

from sympy import integer_nthroot
from RSAExploits.exploits.exploit import Exploit
from RSAExploits.util import compute_priv_rsa_exp

class Fermat(Exploit):
    """ Class providing a run interface to a fermat facorization function"""
    
    def run(self, rsadata_list, info_dict = None):
        """ Factors moduli using fermat factorization 

            This algorithm is best suited for finding prime factors when
            the difference between them is small
    
            Args:
                rsadata_list: A list of RSA_Data objects with moduli to factor
                info_dict: User can specify a bound on the runtime using the
                           key "Fermat" and specifying a number

            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 was factored
        """
        
        print ("Fermat: Running Attack...")
        success = False
        limit = None
        if info_dict != None and "Fermat" in info_dict:
            limit = info_dict["Fermat"]

        for rsadata in rsadata_list:
            if limit != None:
                roots = Fermat.fermat_factor(rsadata.get_n(), limit)
            else:
                roots = Fermat.fermat_factor(rsadata.get_n())

            if roots != None:
                rsadata.set_d(compute_priv_rsa_exp(rsadata.get_e(), roots))
                success = True
        
        if success:     
            print("Fermat: Success, roots found, assigned in rsadata.")
        else:           
            print("Fermat: Failure, roots not found.")
        return success  

    
    @staticmethod   
    def fermat_factor(n, limit = 1000000):
        
        """ Attempt to factor n using fermat factorization
        
            Args:
                n: The modulus to factor
                limit: Limits the runtime of this function

            Assumptions:
                n is the product of 2 primes p and q

            Return:
                The two prime factors p and q
        """             

        a, exact = integer_nthroot(n, 2)
        _max = a + limit
        while a < _max:
            b2 = a*a - n
            if b2 >= 0:
                b, exact = integer_nthroot(b2, 2)
                if b != None:
                    if b*b == b2:
                        break
            a += 1
        if a < _max:
            p = a+b
            q = a-b
            return p, q
        else:
            return None

        


