Diffie–Hellman key exchange

dh.py

# Diffie-Hellman key exchange by Valentino Giudice
# This is just a toy exmaple. It's not ment for security purposes.

from random import randint

def get_prime(bits):
	def miller_rabin(m, a):
		d = 1
		for j in range(bits - 1, -1, -1):
			x = d
			d = (d * d) % m
			if d == 1 and x != 1 and x != m - 1:
				return True
			if (m >> j) % 2 == 0:
				d = (d * a) % m
		return d != 1
	prime = False
	while not prime:
		m = randint(5, 2 ** bits)
		if m % 2 == 0:
			m -= 1
		prime = True
		for i in range(20):
			a = randint(3, m - 1)
			if miller_rabin(m, a):
				prime = False
				break
	return m

def factor(n):
	f = []
	if n % 2 == 0:
		f.append(2)
		while n % 2 == 0:
			n /= 2
	i = 3
	while i <= n:
		if n % i == 0:
			f.append(i)
			while n % i == 0:
				n /= i
		i += 2
	return f

def get_root_key(q, factors):
   is_root = False
   while not is_root:
	   a = randint(2, q - 1)
	   is_root = True
	   for f in factors:
		   if a ** ((q - 1) / f) == 1:
			   is_root = False
			   break
   return a

def expmod_recursive(a, b, q):
	if b == 0:
		return 1
	if b % 2 == 0:
		return (expmod_recursive(a, b / 2, q) ** 2) % q
	return (a * expmod_recursive(a, (b - 1), q)) % q

def expmod_iterative(a, b, q):
	result = 1
	for i in range(b.bit_length(), -1, -1):
		result **= 2
		if (b >> i) % 2 == 1:
			result *= a
		result %= q
	return result

expmod = expmod_iterative

def get_keys(q, a):
	sa = randint(1, q - 1)
	return sa, expmod(a, sa, q)

def get_shared_key(q, sa, pb):
	return expmod(pb, sa, q)

k = 7
q = get_prime(k)
print("q:", q)
f = factor(q - 1)
print("f:", f)
a = get_root_key(q, f)
print("a:", a)
sa, pa = get_keys(q, a)
print("sa:", sa)
print("pa:", pa)
sb, pb = get_keys(q, a)
print("sb:", sb)
print("pb:", pb)
ka = get_shared_key(q, sa, pb)
print("ka:", ka)
kb = get_shared_key(q, sb, pa)
print("kb:", kb)