ECDH implementation in python

Question

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()

Answer 1

• 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.

• Aemyl: 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!

Answer 2

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 :)