example.py

#!/usr/bin/env python2.7

from Crypto.Util.number import bytes_to_long, inverse, long_to_bytes
from Crypto.Random.random import randint

class PublicKey:
	def __init__(self, h, p, g, q):
		self.h = h
		self.p = p
		self.g = g
		self.q = q

class PrivateKey:
	def __init__(self, x, p, g, q):
		self.x = x
		self.p = p
		self.g = g
		self.q = q

def _generate_key():
	"""
	Generate private-public key pair.
	For security reasons, either p should be a safe prime or g should have a
	prime subgroup order. Otherwise it is vulnerable to Short Subgroup Attack.

	:Parameters: _None_

	:Variables:
		g : int/long
			Base point for modular exponentiation.
		p : int/long
			Modulus for modular exponentiation. Should be a safe prime.
		x : int/long
			Receiver's private key, should be kept secret.
		h : int/long
			Receiver's public key
		q : int/long
			Order of group generated by p and equals p-1

	:Return: A tuple containing a Public Key object (class `PublicKey`) and
	a Private Key object (class `PrivateKey`)
	"""
	# Assigning the largest 1024-bit safe prime as p
	p = (1 << 1024) - 1093337
	x = randint(2, p-2)
	g = 7
	q = p - 1
	h = pow(g, x, p)
	pubkey = PublicKey(h, p, g, q)
	privkey = PrivateKey(x, p, g, q)
	return (pubkey, privkey)


def _derive_key(pubkey, x):
    """
    Derive shared secret from the public key of the sender and private key of the
    user.

    :parameters:
        pubkey : object of class `PublicKey`
        x      : secret key of the user

    Returns the derived shared key
    """
    p = pubkey.p
    g = pubkey.g
    h = pubkey.h
    q = pubkey.q
    # h = g^x_alice % p
    # _key = h^x_bob % p = g^(x_alice*x_bob) % p
    _key = pow(h, x, p)
    return _key

if __name__ == "__main__":
    try:
        for i in range(100):
            pubkey, privkey = _generate_key()
            p = pubkey.p
            g = pubkey.g
            x = privkey.x
            y = randint(2, p-2)
            assert pow(g, x*y, p) == _derive_key(pubkey, y)
    except:
        print "[-] Something's wrong in your code! Check the implementation"