16
\$\begingroup\$

I wrote an implementation of elliptic curve Diffie-Hellman key exchange in python. I would like to know if I could make it faster, cleaner or more secure:

def ext_euclidean(a, b):

    t = u = 1
    s = v = 0
    while b:
        a, (q, b) = b, divmod(a, b)
        u, s = s, u - q * s
        v, t = t, v - q * t
    return a, u, v


class Point:

    def __init__(self, x, y):

        self.x = x
        self.y = y

    def __str__(self):

        return '(' + str(self.x) + ', ' + str(self.y) + ')'

    def __eq__(self, p):

        if type(self) != type(p):
            return False
        if self.x == p.x and self.y == p.y:
            return True
        return False


class EllipticCurve:

    def __init__(self, a, b, m):

        self.a = a
        self.b = b
        self.m = m

    def generate(self, x):
        """generate point with given x coordinate."""
        v = (x ** 3 + self.a * x + self.b) % self.m
        legendre = pow(v, self.m // 2, self.m)
        if legendre == (-1) % self.m:
            print("No solution!")
            return None
        elif legendre == 0:
            y = 0
        elif legendre == 1:
            if self.m % 4 == 3:
                y = pow(v, (self.m + 1) // 4, self.m)
                if y ** 2 % self.m != v:
                    print("Algorithm seems to be wrong!")
            elif self.m % 4 == 1:
                print("Algorithm not implemented yet!")
                return None
            else:
                print("Unknown error occured!")
                return None
        else:
            print("Couldn't calculate the Legendre symbol (probably the order of this curve isn't prime)")
            return None
        return Point(x, y)

    def is_generator(self, p):
        """Test if p generates the whole curve or just a sub-group."""
        bm = bin(self.m)[2:]
        if bm == '1' * len(bm):
            # double point until it's None
            # the number of double operations has to be len(bm)
            # only works if self.m is a mersenne prime
            steps = 0
            while p != None:
                p = self.add(p, p)
                steps += 1
            if steps == len(bm):
                return True
            elif steps > len(bm):
                print("Doubled too often!?")
            else:
                print("This point doesn't generate the whole curve!")
            return False
        else:
            return True  # prevent crash

    def add(self, p, q):
        """Point addition on this curve.
        None is the neutral element."""
        if p == None:
            return q
        if q == None:
            return p
        numerator = (q.y - p.y) % self.m
        denominator = (q.x - p.x) % self.m
        if denominator == 0:
            if p == q:
                if p.y == 0:
                    return None
                inv = ext_euclidean(2 * p.y, self.m)[1]
                slope = inv * (3 * p.x ** 2 + self.a) % self.m
            else:
                return None  # p == -q
        else:
            inv = ext_euclidean(denominator, self.m)[1]
            slope = inv * numerator % self.m
        xr = (slope ** 2 - (p.x + q.x)) % self.m
        yr = (slope * (p.x - xr) - p.y) % self.m
        return Point(xr, yr)

    def mul(self, p, n):
        """Binary multiplication.
        Double and add instead of square and multiply."""
        r = None
        for b in bin(n)[2:]:
            r = self.add(r, r)
            if b == '1':
                r = self.add(r, p)
        return r


def main():

    m = 2 ** 521 - 1
    ec = EllipticCurve(1, 0, m)
    print("Curve order:", m)
    while True:
        try:
            x = int(input("x coordinate of generator: ")) % m
            gen = ec.generate(x)
            if gen != None and ec.is_generator(gen):
                break
        except ValueError:
            print("This has to be a number!")
    print("The generator is:", gen)
    while True:
        try:
            a = int(input("Enter private key a: ")) % m
            break
        except ValueError:
            print("This has to be a number!")
    A = ec.mul(gen, a)
    print("Public key A:", A)
    while True:
        try:
            b = int(input("Enter private key b: ")) % m
            break
        except ValueError:
            print("This has to be a number!")
    B = ec.mul(gen, b)
    print("Public key B:", B)
    Ka = ec.mul(B, a)
    print("Ka:", Ka)
    Kb = ec.mul(A, b)
    print("Kb:", Kb)
    if Ka == Kb:
        print("Key generation was successful!")
    else:
        print("Key generation failed!")


if __name__ == '__main__':

    main()
\$\endgroup\$
9
\$\begingroup\$
  • You can easily replace your Point class with namedtuple here. In fact the docs itself contain an example related to a Point class. A namedtuple class comes with a decent __repr__ and comparison will work too.

  • Use better meaningful variable names than u, v, t, m etc.

  • Use is for comparing against singleton None instead of ==.
  • Keep your try-except blocks only around the statement where you're expecting the error, keeping it around multiple statements could make it hard to debug later on. For example check the first try-except in your main() function, it should only be around the int() call.
\$\endgroup\$
  • 1
    \$\begingroup\$ to the variable names: I used u, v, s and t because this letters are also used in the wikipedia description of the extended euclidean algorithm. m could be named module or similar but I don't want to get confused with modules you can import. Good answer though, thanks! \$\endgroup\$ – Aemyl Aug 28 '17 at 4:46
7
\$\begingroup\$
  1. You can write constructs like these:

    if self.x == p.x and self.y == p.y:
        return True
    return False
    

    as:

    return self.x == p.x and self.y == p.y
    
  2. You should avoid interspersing a general algorithm implementation with print statements. Use the python logging framework instead.

  3. There are several checks in the implementation which would warrant to raise an exception instead of printing a message and just returning a default value or or worse simply continuing.

    These would probably warrant raising a ValueError:

    • print("Algorithm seems to be wrong!") in generate, this is particular bad since the computation just moves on and returns a potentially bogus Point.
    • print("Unknown error occured!") in generate
    • print("Doubled too often!?") in is_generator


    This one could raise a NotImplementedError:

    • print("Algorithm not implemented yet!") in generate
  4. You re-compute bm = bin(self.m)[2:] every time is_generator is called. Since m should be considered an invariant you could do this just once in the constructor

  5. Hidden in the comments inside the implementation of is_generator is the requirement that the passed in modulus m has to be a Mersenne prime. This should be at least documented in the constructor invariants.

    Since there is a reasonably efficient test for Mersenne primes in form of the Lucas-Lehmer test you could test this in the constructor. Now obviously if the passed in prime is reasonably large then this could prove impractical but I suspect your implementation would run into other issues in that case as well.

    OTOH there are only 49 Mersenne primes known, so you could store the exponent of all of them and test if n (based on (2^n) - 1) is in the known Mersenne exponent list. Means you'd have to update the code if they find new ones but this seems to take quite some time these days :)

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.