#!/usr/bin/python2
""" Class defintion for the boneh durfee low private exponent exploit
    Much of this file was *autogenerated* from the file boneh_durfee.sage
    taken from: https://github.com/mimoo/RSA-and-LLL-attacks
"""


sage_present = True
import os
path = "/usr/share/sagemath"
os.environ['SAGE_ROOT'] = path
os.environ['SAGE_SRC'] = path
os.environ['SAGE_DOC_SRC'] = path
os.environ['SAGE_LOCAL'] = path
os.environ['DOT_SAGE'] = path
from sage.all_cmdline import *   # import sage library
_sage_const_2 = Integer(2)
_sage_const_1 = Integer(1)
_sage_const_0 = Integer(0)
_sage_const_60 = Integer(60)

# Currently this exploit is set to work if the private
# exponent is <= N ^ 0.27. Theoretical it could be set
# to work as long as d < N ^ 0.292 but that takes a lot
# longer to run
_sage_m_value = Integer(6)
_sage_delta_value = RealNumber('0.27')


debug = False

class Boneh_Durfee():
    """ Class providing a run interface to the boneh_durfee exploit"""

    def run(self, n, e, info_dict = None):

        """ Attempts to recover private exponent using the boneh durfee exploit
            This attack works when the private key d is < N ^ 0.27. This limit
            can be tweaked up to a theoretical N^0.292, but the computation
            begins to take a very long time.
        """

        global sage_present

        if not sage_present:
            print("Boneh Durfee: Sage required, skipping exploit")
            return

        # print("Boneh Durfee: Running Attack...")
        success = False

        solx, soly = Boneh_Durfee.boneh_durfee_wrapper(n, e)
        if solx == None:
            exit(1)

        A = int((n+_sage_const_1 )/_sage_const_2 )
        d = (solx * (A + soly) + _sage_const_1) / e
        success = True
        d = (int(d))

        if success:
            # print("Boneh Durfee: Success, at least one private exponent recovered")
            print ("DFOUND:")
            print (d)
            print ("#")
            return d
        else:
            print("Boneh Durfee: Failure, no private exponents recovered")
            exit(2)



    @staticmethod
    def boneh_durfee_wrapper(n, e):
        """ Wrapper for the boneh durfee exploit """
        P = PolynomialRing(Zmod(e), names=('x', 'y',)); (x, y,) = P._first_ngens(2)
        A = int((n+_sage_const_1 )/_sage_const_2 )
        pol = _sage_const_1  + x * (A + y)

        delta = _sage_delta_value
        m = _sage_m_value
        t = int((_sage_const_1 - _sage_const_2*delta)*m)

        X = _sage_const_2*floor(n**delta)
        Y = floor(n**(_sage_const_1 / _sage_const_2))

        return Boneh_Durfee.boneh_durfee(pol, e, m, t, X, Y)


    @staticmethod
    def boneh_durfee(pol, modulus, mm, tt, XX, YY):
        """
        Boneh and Durfee revisited by Herrmann and May

        finds a solution if:
        * d < N^delta
        * |x| < e^delta
        * |y| < e^0.5
        whenever delta < 1 - sqrt(2)/2 ~ 0.292
        """

        # substitution (Herrman and May)
        PR = PolynomialRing(ZZ, names=('u', 'x', 'y',)); (u, x, y,) = PR._first_ngens(3)
        Q = PR.quotient(x*y + _sage_const_1  - u) # u = x*y + 1
        polZ = Q(pol).lift()

        UU = XX*YY + _sage_const_1

        # x-shifts
        gg = []

        for kk in range(mm + _sage_const_1 ):
            for ii in range(mm - kk + _sage_const_1 ):
                xshift = x**ii * modulus**(mm - kk) * polZ(u, x, y)**kk
                gg.append(xshift)

        gg.sort()

        # x-shifts monomials
        monomials = []

        for polynomial in gg:
            for monomial in polynomial.monomials():
                if monomial not in monomials:
                    monomials.append(monomial)

        monomials.sort()

        # y-shifts (selected by Herrman and May)
        for jj in range(_sage_const_1 , tt + _sage_const_1 ):
            for kk in range(floor(mm/tt) * jj, mm + _sage_const_1 ):
                yshift = y**jj * polZ(u, x, y)**kk * modulus**(mm - kk)
                yshift = Q(yshift).lift()
                gg.append(yshift) # substitution

        # y-shifts monomials
        for jj in range(_sage_const_1 , tt + _sage_const_1 ):
            for kk in range(floor(mm/tt) * jj, mm + _sage_const_1 ):
                monomials.append(u**kk * y**jj)

        # construct lattice B
        nn = len(monomials)

        BB = Matrix(ZZ, nn)

        for ii in range(nn):

            BB[ii, _sage_const_0 ] = gg[ii](_sage_const_0 , _sage_const_0 , _sage_const_0 )

            for jj in range(_sage_const_1 , ii + _sage_const_1 ):
                if monomials[jj] in gg[ii].monomials():
                    BB[ii, jj] = gg[ii].monomial_coefficient(monomials[jj]) * monomials[jj](UU,XX,YY)

        # check if vectors are helpful
        if debug:
            helpful_vectors(BB, modulus**mm)

        # check if determinant is correctly bounded
        if debug:
            det = BB.det()
            bound = modulus**(mm*nn)
            if det >= bound:
                # print("We do not have det < bound. Solutions might not be found.")
                diff = (log(det) - log(bound)) / log(_sage_const_2 )
                # print("size det(L) - size e^(m*n) = ", floor(diff))
            else:
                print("det(L) < e^(m*n)")

        # debug: display matrix
        if debug:
            matrix_overview(BB, modulus**mm)

        # LLL
        BB = BB.LLL()

        # vector 1 & 2 -> polynomials 1 & 2
        PR = PolynomialRing(ZZ, names=('w', 'z',)); (w, z,) = PR._first_ngens(2)

        pol1 = pol2 = _sage_const_0
        for jj in range(nn):
            pol1 += monomials[jj](w*z+_sage_const_1 ,w,z) * BB[_sage_const_0 , jj] / monomials[jj](UU,XX,YY)
            pol2 += monomials[jj](w*z+_sage_const_1 ,w,z) * BB[_sage_const_1 , jj] / monomials[jj](UU,XX,YY)

        # resultant
        PR = PolynomialRing(ZZ, names=('q',)); (q,) = PR._first_ngens(1)
        rr = pol1.resultant(pol2)

        if rr.is_zero() or rr.monomials() == [_sage_const_1 ]:
            if debug:
                print("failure")
            return None, None

        rr = rr(q, q)
        # solutions
        try:
            soly = rr.roots()[_sage_const_0 ][_sage_const_0 ]
            if debug:
                print("found for y_0:", soly)

            ss = pol1(q, soly)
            solx = ss.roots()[_sage_const_0 ][_sage_const_0 ]
            if debug:
                print("found for x_0:", solx)
        except IndexError:
            return None, None
        return solx, soly



# display stats on helpful vectors
def helpful_vectors(BB, modulus):
    nothelpful = _sage_const_0
    for ii in range(BB.dimensions()[_sage_const_0 ]):
        if BB[ii,ii] >= modulus:
            nothelpful += _sage_const_1

    # print(nothelpful, "/", BB.dimensions()[_sage_const_0 ], " vectors are not helpful")

# display matrix picture with 0 and X
def matrix_overview(BB, bound):
    for ii in range(BB.dimensions()[_sage_const_0 ]):
        a = ('%02d ' % ii)
        for jj in range(BB.dimensions()[_sage_const_1 ]):
            a += '0' if BB[ii,jj] == _sage_const_0  else 'X'
            if BB.dimensions()[_sage_const_0 ] < _sage_const_60 :
                a += ' '
        if BB[ii, ii] >= bound:
            a += '~'
        # print(a)



import sys
Boneh_Durfee().run(int(sys.argv[1]), int(sys.argv[2]))
# crack(137709853388874260067664060306224801065880567280896538344708662296491187670268243771581008615887329285442080267493172576517227603581814922784005894668293664250883858847270414696279127204843462302026428548409668060539273696154323753582118715764251826181439387752900169898505518594134423187783089073450520527969,85605944479801539190292577400294315989703364683792339737453887619652946907819548388390076584246192710348319524405250340541547224543819349494488596679810310694697763188201375896240755682682520811114377302206659716598003967368992125832908158096663258971132148617761060030524535044970322206572134828548118438019)
