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 _encrypt(message, pubkey):
	"""
	Encrypt message using ElGamal encryption system

	:Parameters:
		message : str
				  plaintext to be encrypted
		pubkey  : instance of `PublicKey` class
				  Alice's public key parameters used for encryption

	:Variables:
		g : int/long
			Base point for modular exponentiation.
		p : int/long
			Modulus for modular exponentiation. Should be a safe prime.
		h : int/long
			Receiver's public key
		q : int/long
			Order of group generated by p and equals p-1
		y : int/long
			Ephemeral key generated by the sender
		c1, c2: int/long
			Ciphertext pair
		s : int/long
			Shared secret

	:Return:
		A tuple containing ciphertext pair c1, c2
	"""
	h = pubkey.h
	p = pubkey.p
	g = pubkey.g
	q = pubkey.q
	m = bytes_to_long(message)
	# Generating ephemeral key: `y`
	y = randint(2, p-2)
	c1 = pow(g, y, p)
	# Computing the shared secret: `s`
	s = pow(h, y, p)
	c2 = (m*s) % p
	return (c1, c2)

def _decrypt(ciphertext, privkey):
	"""
	Decrypt ciphertext using ElGamal Encryption System

	:Parameters:
		ciphertext : int/long tuple
			Ciphertext of ElGamal encrypted plaintext
		privkey : instance of `PrivateKey` class
		 	Receiver's private key used for decryption

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

	"""
	c1, c2 = ciphertext
	g = privkey.g
	p = privkey.p
	x = privkey.x
	s = pow(c1, x, p)
	m = (c2*inverse(s, p)) % p
	return m

if __name__ == "__main__":
	from os import urandom
	pubkey, privkey = _generate_key()
	for i in range(100):
		message = chr(1) + urandom(16)
		ct = _encrypt(message, pubkey)
		try:
	 		assert long_to_bytes(_decrypt(ct, privkey)) == message
		except:
			print "[-] Something's wrong! Check the implementation!"
			import sys
			sys.exit()