"""
Proof-of-concept code for some attacks described in 'Attacking
Cryptographic Schemes Based on "Perturbation Polynomials"' by Martin
Albrecht, Craig Gentry, Shai Halevi and Jonathan Katz.

AUTHOR: 
    -- 2009-01 Martin Albrecht <M.R.Albrecht@rhul.ac.uk>: initial version
"""

class MobiHoc2007:
    """
    Class for experimenting with the key predistribution scheme from
 
    W. Zhang, M. Tran, S. Zhu, and G. Cao. A Random Perturbation-based
    Scheme for Pairwise Key Establishment in Sensor Networks. 8th ACM
    International Symposium on Mobile Ad Hoc Networking and Computing
    (MobiHoc'07), 2007.
    """
    def __init__(self, p, t, r):
        """
        Create new MobiHoc2007 instance:

        INPUT:
            p -- a prime
            t -- a degree > 0
            r -- an integer < p
        """
        self.p = p
        self.t = t
        self.r = r
        self.N = t + 3
        self.u = 2

    def random_symmetric_polynomial(self):
        """
        Returns and stores a random symmetric polynomial of degree t
        over GF(p)[x,y].
        """
        verbose("",level=1)
        p = self.p
        t = self.t
        P = PolynomialRing(GF(p), 2, 'x,y')
        x,y = P.gens()

        _, total = P._precomp_counts(2, t)
        FF = P.random_element(degree=t, terms=total)

        F = 0
        for c,m in FF:
            a,b = m.exponents()[0]
            F += c* x**a *y**b + c* y**a * x**b
        self.F = F
        return F

    def random_perturbation_polynomials(self):
        """
        Returns and stores two random perturbation polynomials of
        degree t over GF(p)[y]. Also a set S is returned and stored
        which contains all elements in GF(p) for which the two
        polynomials evaluate to values < r.
        """
        _cpu = cputime()
        verbose("",level=1)
        p = self.p
        t = self.t
        r = self.r
        N = self.N
        
        K = GF(p)
        P = PolynomialRing(K, 'y')
        x, = P.gens()

        g = P.random_element(degree=t)
        h = P.random_element(degree=t)
        S = []
        i = 0
        step = int(p)/10

        try:
            # Try if we have our hack applied which implements
            # evaluation of all points on C level avoiding Python
            # overhead. Fall back to Python if that fails.
            S = g.eval_on_all_points(h, r)
            S = map(K,S)
        except AttributeError:
            e = K(0)
            end = K(p-1)
            o = K(1)

            while e!=end:
                if i==step:
                    print ".",
                    step+=int(p)/10
                    sys.stdout.flush()
                i+=1
                if g(e) < r and h(e) < r:
                    S.append(e)
                e+=o
            print

        self.g = g
        self.h = h
        self.S = S
        print "Setup time %8.3f"%cputime(_cpu)
        return g,h,S

    def __getattr__(self, name):
        if name in ('g', 'h', 'S'):
            _ = self.random_perturbation_polynomials()
        elif name == 'F':
            _ = self.random_symmetric_polynomial()
        else:
            raise AttributeError(name)
        return getattr(self, name)

    def random_node(self, not_in=None):
        """
        Return a random node which satisfies the random perturbation
        polynomials.
        """
        verbose("",level=1)
        if not_in is None:
            not_in = set()
        else:
            not_in = [n.x for n in not_in]
        S = list(set(self.S).difference(not_in))

        i = randint(0,len(S)-1)
        return self.node(self.S.index(S[i]))

    @cached_method
    def node(self, i):
        """
        Return the i-th node from the set of nodes which satisfy the
        perturbation polynomials.

        The result is cached.

        INPUT:
            i -- an integer >= 0 and < len(S)
        """
        g, h, S, F = self.g, self.h, self.S, self.F
        r = self.r
        ai = randint(-self.u,self.u)
        bi = randint(-self.u,self.u)
        assert(g(S[i]) < r)
        assert(h(S[i]) < r)
        _,y = F.parent().gens()
        fi  =F(S[i],y).univariate_polynomial()
        Si  = fi + ai*g + bi*h
        return MH07Node(S[i], fi, Si, ceil(log_b(self.r,2)))

    def shared_secret(self, i=None, j=None):
        """
        Return the shared secret between the nodes indexed by i and j
        in the set S.

        INPUT:
            i -- an integer
            j -- an integer
        """
        if i is None:
            i = self.random_node()
        if j is None:
            j = self.random_node()
        return i.shared_secret(j)

    def verify_noise_space(self, S):
        """
        Given a matrix S with two rows verify that the space spanned
        by these two rows is the same as the space spanned by the
        perturbation polynomial coefficient vectors.

        INPUT:
            S -- a matrix over GF(p) with two rows.
        """
        assert(S.rank() == 2)
        t = self.t
        M = Matrix(GF(self.p),2, t + 1)
        for i in range(t + 1):
            M[0,i] = self.g[i]
        for i in range(t + 1):
            M[1,i] = self.h[i]
        M = M.stack(S)
        return M.rank() == 2

    @cached_method
    def noise_space(self, small=True):
        """
        Identify the spaces spanned by the coefficient vectors of the
        perturbation polynomials.

        The result is not verified.

        INPUT:
            small -- if True LLL is run to recover 'small' candidates
                     for basis vectors (default: True)

        OUTPUT:
            a matrix with two rows representing the two vectors
            spanning the space.
        """
        verbose("",level=1)
        t = self.t
        N = self.N
        nodes = set()
        while (len(nodes) < N):
            n = self.random_node(nodes)
            nodes.add(n)
        nodes = list(nodes)
        verbose("done",level=1)
        self.compromised_nodes = nodes

        X = [e.x for e in nodes]
        S = Matrix(GF(self.p), N, t+1)
        for i in range(N):
            for j in range(t+1):
                S[i,j] = nodes[i].S[j]

        for i in range(N-2):
            S[N-2] -= L(X[:N-2],X[N-2],i) * S[i] 

        for i in range(N-2):
            S[N-1] -= L(X[:N-2],X[N-1],i) * S[i] 

        if small:
            return self.small_linear_combinations(S.matrix_from_rows([N-2,N-1]))
        else:
            return S.matrix_from_rows([N-2,N-1])

    def get_LLL(self, NS, n):
        """
        Return the LLL reduced matrix B which was used to recover
        'small' basis vectors for NS. The dimension of B will be 
        (n+2) x n.

        INPUT:
            NS -- a matrix with two rows
            n -- an integer > 0
        """
        node = self.compromised_nodes[0]
        P = node.S.parent()
        v = P(NS[0].list())
        w = P(NS[1].list())

        p, t, compromised_nodes = self.p, self.t, self.compromised_nodes

        nodes = set()
        while len(nodes) < n:
            nodes.add(self.random_node(nodes))

        nodes = list(nodes)

        vals = [e.x for e in nodes]
        A = Matrix(ZZ,n+2,n)

        for i in range(n):
            xi = vals[i]
            A[0,i] = v(xi)
            A[1,i] = w(xi)
            A[i+2,i] = p
        B = A.LLL()

        return B, vals

    def small_linear_combinations(self, NS):
        """
        Given the noise space spanned by the two rows of NS, return
        another matrix A where the two rows are 'small' basis vectors
        also spanning NS.

        INPUT:
            NS -- a matrix with two rows.
        """
        n = self.compromised_nodes[0]
        r = self.r
        P = n.S.parent()
        v = P(NS[0].list())
        w = P(NS[1].list())

        p, t, compromised_nodes = self.p, self.t, self.compromised_nodes

        n = 40

        nodes = set()
        while len(nodes) < n:
            nodes.add(self.random_node(nodes))

        nodes = list(nodes)

        vals = [e.x for e in nodes]
        A = Matrix(ZZ,n+2,n)

        for i in range(n):
            xi = vals[i]
            A[0,i] = v(xi)
            A[1,i] = w(xi)
            A[i+2,i] = p
        B = A.LLL()

        # now solve for a,b,c,d
        M = Matrix(GF(p),4,5)
        M[0,0] = A[0,0]
        M[0,1] = A[1,0]
        M[0,4] = -B[2,0]
        M[1,0] = A[0,1]
        M[1,1] = A[1,1]
        M[1,4] = -B[2,1]

        M[2,2] = A[0,0]
        M[2,3] = A[1,0]
        M[2,4] = -B[3,0]
        M[3,2] = A[0,1]
        M[3,3] = A[1,1]
        M[3,4] = -B[3,1]

        E = M.echelon_form()

        a,b,c,d = [-each for each in list(E.column(4))]
        g = a*v + b*w
        h = c*v + d*w

        S = set([g,-g,h,-h,g+h,g-h,-g+h,-g-h])
        polys = []
        while len(S) > 0:
            f = S.pop()
            for n in self.compromised_nodes:
                if ZZ(f(n.x)) >= r:
                    break
            else:
                polys.append(f)
        assert(len(polys) == 2)
        g,h = polys

        return Matrix(GF(p), 2, t+1, g.list() + h.list())
        
    def fi_equation_system(self):
        r"""
        Return an equation system to recover the $f_i$'s.
        """
        p, t, r = self.p, self.t, self.r
        N = self.N
        VS = VectorSpace(GF(p), t+1)

        NS = self.noise_space()

        assert(self.verify_noise_space(NS))

        nodes = self.compromised_nodes

        P = PolynomialRing(GF(p), 2*N, ["gamma%d"%i for i in range(N)]+["delta%d"%i for i in range(N)])
        gamma = P.gens()[:N]
        delta = P.gens()[N:]

        v = VS(NS[0])
        w = VS(NS[1])
        X = [e.x for e in nodes]

        equations = []

        for i in range(N):
            xi = VS([nodes[i].x**exp for exp in xrange(t+1)])
            Si = nodes[i].S.list()
            assert(len(Si) == t+1)
            Si = VS(Si)
            for j in range(i+1, N):
                xj = VS([nodes[j].x**exp for exp in xrange(t+1)])
                Sj = nodes[j].S.list()
                assert(len(Sj) == t+1)
                Sj = VS(Sj)
                lhs = Si.dot_product(xj) - gamma[i] * v.dot_product(xj) - delta[i] * w.dot_product(xj)
                rhs = Sj.dot_product(xi) - gamma[j] * v.dot_product(xi) - delta[j] * w.dot_product(xi)
                equations.append( lhs - rhs )

        St1 = VS(nodes[N-2].S.list())
        for j in range(N):
            xj = VS([nodes[j].x**exp for exp in xrange(t+1)])
            ft1 = St1.dot_product(xj) - gamma[N-2] * v.dot_product(xj) - delta[N-2] * w.dot_product(xj)
            for i in range(N-2):
                Si = VS(nodes[i].S.list())
                mult = L(X[:N-2],X[N-2],i) 
                fi = Si.dot_product(xj) - gamma[i] * v.dot_product(xj) - delta[i] * w.dot_product(xj)
                ft1 -= mult *fi
            equations.append(ft1)

        St2 = VS(nodes[N-1].S.list())
        for j in range(N):
            xj = VS([nodes[j].x**exp for exp in xrange(t+1)])
            ft2 = St2.dot_product(xj) - gamma[N-1] * v.dot_product(xj) - delta[N-1] * w.dot_product(xj)
            for i in range(N-2):
                Si = VS(nodes[i].S.list())
                mult = L(X[:N-2],X[N-1],i) 
                fi = Si.dot_product(xj) - gamma[i] * v.dot_product(xj) - delta[i] * w.dot_product(xj)
                ft2 -= mult *fi
            equations.append(ft2)


        return mq.MPolynomialSystem(P,equations)

    def guess_and_verify(self, MQ, a, b, c, pick_randomly):
        r"""

        Subsititute $a,b,c$ for
        $\gamma_{t+1},\delta_{t+1},$\gamma_{t+2}$ respectively and
        verify if all other $\delta$'s and $\gamma$'s are 'small'
        under this guess. Small in this context means in [-u,u].

        INPUT:
            MQ -- an equation system
            a -- an element in GF(p)
            b -- an element in GF(p)
            c -- an element in GF(p)
            pick_randomly -- pick some sufficiently large random
                             subset of equations instead of
                             considering all equations.
        """
        p = self.p
        N = self.N
        t = self.t
        MQ = copy(MQ)
        K = GF(p)
        P = MQ.ring()
        gamma = P.gens()[:N]
        delta = P.gens()[N:]

        if pick_randomly:
            MQ2 = [gamma[-2] - a, delta[-2] -b, gamma[-1] - c]

            for i in range(6*N):
                MQ2.append(MQ[randint(0,MQ.ngens()-1)])

            MQ = mq.MPolynomialSystem(MQ2)
        else:
            MQ += [gamma[-2] - a, delta[-2] -b, gamma[-1] - c]

        V = self.linear_variety(MQ)
        assert(len(V) == 1)
        V = V[0]
        V[gamma[-2]] = a
        V[delta[-2]] = b
        V[gamma[-1]] = c

        for e in gamma[:t]:
            if abs(ZZ(V[e])) > self.u and abs(ZZ(V[e])-p) > self.u:
                return None
        for e in delta[:t]:
            if abs(ZZ(V[e])) > self.u and abs(ZZ(V[e])-p) > self.u:
                return None

        return V, gamma, delta

    def linear_variety(self, MQ):
        r"""
        Return the variety for \var{MQ}
        
        INPUT:
            MQ -- an equation system.
        """
        A,v = MQ.coefficient_matrix(sparse=False)
        E = A.echelon_form()
        assert(A.rank() == A.ncols() - 1)
        F = [f for f in (E*v).list()]
        V = {}
        for f in F:
            if f==0:
                continue
            V[f.variables()[0]] = -f.univariate_polynomial().constant_coefficient()
        return [V]

    def recover_master_secret(self):
        """
        Run the attack on this instance.
        """
        p = self.p
        t = self.t
        r = self.r
        N = self.N

        print "SYSTEM PARAMETERS"
        print "p",p,"t",t,"r",r

        _ = self.random_perturbation_polynomials()

        print "ATTACK VERIFICATION"

        K = GF(p)

        _cpu = cputime()
        NS = self.noise_space()
        print "Noise Space time: %8.3f"%(cputime(_cpu))
        v = K['y'](NS[0].list())
        w = K['y'](NS[1].list())

        print "Rank of the noise space:",(NS.rank())

        A = Matrix(K,2, t+1,self.g.list() + self.h.list())

        print "Noise spaces match:", A.echelon_form() == NS.echelon_form()

        print r"v  \in \{g, h\}:", v in [self.g, self.h]
        print r"v' \in \{g, h\}:", w in [self.g, self.h]

        _wall = walltime()
        MQ = self.fi_equation_system()

        V = None
        for a in range(-self.u, self.u+1):
            print a,
            sys.stdout.flush()
            for b in range(-self.u, self.u+1):
                for c in range(-self.u, self.u+1):
                    sol = self.guess_and_verify(MQ,a,b,c, pick_randomly=True)
                    if sol is not None:
                        print "solution found!"
                        break
                if sol is not None:
                    break
            if sol is not None:
                break
        if sol is None:
            raise ArithmeticError("no solution found")

        print "Brute force time %8.3f"%(walltime(_wall))
        V, gamma, delta = sol

        X,Y = self.F.parent().gens()
        myF = []
        R = PolynomialRing(K, t+1, 'c')
        nodes = self.compromised_nodes
        x_coeffs = R.gens()

        for y_deg in xrange(t+1):
            equations = []
            for i in range(len(nodes)):
                fi = nodes[i].S - V[gamma[i]]*v - V[delta[i]]*w
                equations.append(sum([x_coeffs[j]*nodes[i].x**j for j in range(len(x_coeffs))], -fi[y_deg]))
            V_x =  self.linear_variety(mq.MPolynomialSystem(equations))
            assert(len(V_x) == 1)
            V_x = V_x[0]
            myF.append(sum([V_x[gen] * X**i for i,gen in enumerate(x_coeffs)]) * Y**y_deg)

        myF = sum(myF)

        print "F* matches F:",myF == self.F
        return myF == self.F

    def express_basis_in_g_and_h(self, B, vals):
        r"""
        Given a basis as rows of \var{B} express this as linear
        combinations of $g$ and $h$.

        INPUT:
            B -- a matrix containing the basis
            vals -- used to evaluate $g$ and $h$
        """
        p = self.p

        g = [self.g(each) for each in vals]
        h = [self.h(each) for each in vals]

        M = Matrix(GF(p), 2*B.nrows(), 2*B.nrows() + 1)
        ncols = M.ncols()
        for i in xrange(B.nrows()):
            M[2*i  , 2*i    ] = g[0]
            M[2*i  , 2*i+1  ] = h[0]
            M[2*i  , ncols-1] = B[i,0]

            M[2*i+1, 2*i    ] = g[1]
            M[2*i+1, 2*i+1  ] = h[1]
            M[2*i+1, ncols-1] = B[i,1]

        M.echelonize()

        g_coeff = []
        h_coeff = []

        for i in xrange(B.nrows()):
            g_coeff.append(-M[2*i, ncols-1])
            h_coeff.append(-M[2*i+1,ncols-1])

        # some checks
        assert(g_coeff[0] * g[0] + h_coeff[0] * h[0] == B[0,0])
        assert(g_coeff[0] * g[2] + h_coeff[0] * h[2] == B[0,2])

        assert(g_coeff[1] * g[0] + h_coeff[1] * h[0] == B[1,0])
        assert(g_coeff[1] * g[2] + h_coeff[1] * h[2] == B[1,2])

        l = []
        for coeff in g_coeff:
            if ZZ(coeff) > p//2:
                l.append(ZZ(coeff)-p)
            else:
                l.append(ZZ(coeff))
        g_coeff = l

        l = []
        for coeff in h_coeff:
            if ZZ(coeff) > p//2:
                l.append(ZZ(coeff)-p)
            else:
                l.append(ZZ(coeff))
        h_coeff = l

        return g_coeff, h_coeff

    def check_LLL(self, rs, trials=3):
        r"""
        For $e \in rs$ run \var{trials} and check whether the LLL step
        was successful.

        INPUT:
            rs -- a range (e.g. [1,2,3])
            trials -- the number of trials to run for each parameter.
        """
        p = self.p
        t = self.t
        for e in rs:
            r = 2**e
            n = ceil(3*(log(p)/(log(p) - log(2*r))).n())
            sys.stdout.flush()
            alltme = 0.0

            for trial in range(trials):
                mh = MobiHoc2007(p=p, t=t, r=r)
                tme = cputime()
                B,vals = mh.get_LLL(mh.noise_space(small=True), n=n)
                g_coeff, h_coeff = mh.express_basis_in_g_and_h(B, vals)
                assert(abs(g_coeff[2]) <= 1 and abs(g_coeff[3]) <= 1 and abs(h_coeff[2]) <= 1 and abs(h_coeff[3]) <= 1)

                sys.stdout.flush()
                alltme += cputime(tme)

            print "------------------------------------------"
            print "2**%s"%e, n, alltme/trials
            sys.stdout.flush()


class MH07Node:
    """
    A Node in a sensor network capable of executing the protocol from
    MobiHoc 2007.
    """
    def __init__(self, x, f, S, bound):
        self.x = x
        self.f = f
        self.S = S
        self.bound = bound

    def shared_secret(self, other):
        a = ZZ(self.S(other.x))
        b = ZZ(other.S(self.x))
        return a.digits(2)[self.bound:] == b.digits(2)[self.bound:]

    def __cmp__(self, other):
        return cmp(self.x,other.x)
    def __hash__(self):
        return hash((self.x,self.S, self.bound))
    def __repr__(self):
        return "<node %s>"%self.x

def L(X, x, i):
    xi = X[i]
    prod = 1
    for j,xj in enumerate(X):
        if i == j:
            continue
        prod *= (x - xj)/(xi-xj)
    return prod

class PerCom07:
    """
    Class for experimenting with the key predistribution scheme from

    N.V. Subramanian, C. Yang, and W. Zhang. Securing Distributed Data
    Storage and Retrieval in Sensor Networks. 5th IEEE International
    Conference on Pervasive Computing and Communications (PerCom'07),
    2007.
    """
    def __init__(self, p, t, r):
        r"""
        INPUT:
            p -- a prime ($\approx 2^{64}$)
            r -- an integer 0 < r < p ($\approx 2^{16})
            t -- an integer ($\approx 15$)
        """
        self.p = p
        self.t = t
        self.r = r

    def random_master_polynomial(self):
        """
        Return and store a random master polynomial.
        """
        p,t,r = self.p,self.t,self.r
        Pu = PolynomialRing(GF(p),'u')
        Pv = PolynomialRing(GF(p),'v')
        P = PolynomialRing(GF(p),'u,v')
        Fu = Pu.random_element(degree=t)
        Fv = Pv.random_element(degree=t)
        self.F = P(Fu)*P(Fv)
        return self.F

    def __getattr__(self, name):
        if name == 'F':
            _ = self.random_master_polynomial()
        else:
            raise AttributeError(name)
        return getattr(self, name)

    @cached_method
    def phase_polynomial(self, v):
        r"""
        Return a polynomial for the phase \var{v}.
        """
        b_u = randint(0,self.r)
        return self.F.subs(v=v).univariate_polynomial() + b_u

    @cached_method
    def node_polynomial(self, u):
        r"""
        Return a polynomial for the node \var{v}.
        """
        a_u = randint(0,self.r)
        return self.F.subs(u=u).univariate_polynomial() + a_u, a_u
    
    def node(self, u):
        r"""
        Return a new node for u.
        """
        return PC08Node(u,*self.node_polynomial(u))

    def phase(self, v):
        r"""
        Return a new phase for v.
        """
        return PC08Phase(v,self.phase_polynomial(v))

    def random_node(self):
        r"""
        Return a random node.
        """
        return self.node(u=randint(0,self.p-1))

    def random_phase(self):
        r"""
        Return a random phase.
        """
        return self.phase(v=randint(0,self.p-1))

    def key(self, n, v):
        r"""
        Return the key for the node \var{n} for the phase \var{v}.
        """
        return n.f(v)

    def recover_F1(self, attack0=False):
        r"""
        Recover $F_1$.
        """
        t = self.t
        p = self.p

        Q = self.F.parent()
        u,v = Q.gens()

        N = set()
        while len(N) < t+1:
            N.add(self.random_node())

        if attack0:
            P = set()
            while len(P) < t+1:
                P.add(self.random_phase())

        _sum = []
        for deg in range(1,t+1):
            R = PolynomialRing(GF(p),'u')
            X = [(n.u,n.f[deg]) for n in N]
            _sum.append( R.lagrange_polynomial(X)*v**deg )
        S_u =  set(list(iter(sum(_sum))))

        if attack0:
            _sum = []
            for deg in range(1,t+1):
                R = PolynomialRing(GF(p),'v')
                X = [(x.v,x.f[deg]) for x in P]
                _sum.append( R.lagrange_polynomial(X) *u**deg )
            S_v = set(list(iter(sum(_sum))))

        if attack0:
            return sum(map(mul,S_u.intersection(S_v).union(S_u.difference(S_v).union( S_v.difference(S_u) ))))
        else:
            return sum(map(mul,S_u))

    def check_LLL(self, n=None):
        """
        Verify the heuristic in Section 4.2.
        """
        p,t,r = self.p,self.t,self.r
        F1 = self.recover_F1(attack0 = False)
        v = randint(0, self.p-1)

        if n is None:
            n = ceil(((t+1)*log(p)/(log(p)-log(2*r))).n())
        
        N = set()
        while len(N) < n + 1:
            N.add(self.random_node())
        N = list(N)
        
        y = lambda i: N[i].f(v) - F1(N[i].u, v)


        A = Matrix(ZZ,t+1+n+1,n)
        for i in range(n):
            A[0,i] = y(i+1) - y(i)
        for j in range(1,t+1):
            for i in range(n):
                A[j,i] = N[i+1].u**j - N[i].u**j
        for j in range(n):
            A[t+1+j,j] = p
        print A
        B = A.LLL()
        for row in B.rows():
            if not row.is_zero():
                break
       
        return all(abs(row[i]) == abs(N[i+1].a - N[i].a)  for i in range(n))
        

class PC08Node:
    def __init__(self, u, fu, a):
        self.u = u
        self.f = fu
        self.a = a
    def __cmp__(self, other):
        return cmp(self.u,other.u)
    def __hash__(self):
        return hash((self.u,self.f))
    def __repr__(self):
        return "<node %s>"%self.u

class PC08Phase:
    def __init__(self, v, fv):
        self.v = v
        self.f = fv
    def __cmp__(self, other):
        return cmp(self.v,other.v)
    def __hash__(self):
        return hash((self.v,self.f))
    def __repr__(self):
        return "<phase %s>"%self.u