""" Class defintion for Wiener's attack """

import re
import math
from sympy.solvers import solve
from sympy.core.numbers import Integer
from sympy import Symbol
from RSAExploits.exploits.exploit import Exploit
from RSAExploits.util import compute_priv_rsa_exp

class Wiener(Exploit):
    """ Class providing a run interface to Wiener's attack """

    def run(self, rsadata_list, info_dict = None):
        """ Attempts to recover private exponent using Wiener's attack

            This attack works when the private key d is < N ^ 0.25. 
            
            Args:
                rsadata_list: A list of RSA_Data objects on which to attempt
                              Wiener's attack
                info_dict: Not currently used

            Side Effects: 
                If a private key is recovered, the corresponding field in the
                RSA_Data object will be set

            Return: 
                True if at least one private key is recovered
        """
    
        print ("Wiener: Running Attack...")
        success = False

        for rsadata in rsadata_list:
            roots = self.wiener_attack(rsadata.get_n(), rsadata.get_e())
                
            # If wiener's attack was successful for the current rsaData object
            if roots != None:
                rsadata.set_d(compute_priv_rsa_exp(rsadata.get_e(), roots))
                success = True

        if success:     
            print("Wiener: Success, roots found, assigned in rsadata.")
            return True                         
        else:           
            print("Wiener: Failure, roots not found.")
            return False

    @staticmethod
    def conFraction(fraction):
        (a,b) = fraction
        v=[]
        v.append(0)
        while not a == 1:
            (x,y) = fraction
            try:
                r =  y / x
                a = y % x
            except ZeroDivisionError as details:
                print ("Exception occured:",details)
                print("Continuing to next RSA data...")
                return None
            b=x
            fraction = (a,b)
            v.append(r)
        v.append(b)
        return v

    @staticmethod
    def makeIndexedConvergent(sequence, index):
        (a,b)=(1,sequence[index])
        while index>0:
            index-=1
            (a,b)=(b,sequence[index]*b+a)
        return (b,a)

    @staticmethod
    def makeConvergents(sequence):
        r=[]
        for i in xrange(0,len(sequence)):
            r.append(Wiener.makeIndexedConvergent(sequence,i))
        return r

    @staticmethod
    def wiener_attack(N,e):
        v = Wiener.conFraction((e,N))
        if v is None:
            return None
        conv = Wiener.makeConvergents(v)
        for frac in conv:
            (k,d)=frac
            if k == 0:
                continue
            phiN=((e*d)-1)/k
            if ((pow(N-phiN+1, 2) - 4*N) <= 0 ): # (b^2 -4ac) <= 0 
                continue
                
            # solving equation (a**2 + b*x + c)
            x = Symbol('x', integer=True)
            roots = []
            roots = solve(x**2 -(N-phiN+1)*x + N, x, rational=False)
     
            if len(roots) == 2:
                if type(roots[0]) is Integer and type(roots[0]) is Integer:
                             
                    if(roots[0]*roots[1] == N):
                        return roots
         

