import math,fractions, random
from operator import itemgetter

class Ecc:

	def __init__(self,p=5):
		self.id = "Point at infinity"
		self.p = p

	def setCurve(self,a,b):
		self.a = a
		self.b = b

	def setField(self,prime):
		self.p = prime

	def printField(self):
		print self.p

	def checkPoint(self,p):
		if not p == self.id:
			ls = (p[1]*p[1])%self.p
			hs = ((p[0]**3)+(self.a*p[0]) + self.b)%self.p
			if ls == hs:
				return True
			else: 
				return False
		else:
			return True

	# find n*p, where n is the number and p the point
	def multiplumOfPoint(self,p,n):
		if self.checkPoint(p):
			binrep = bin(n)[2:]
			res = self.id
			for bit in binrep:
				res = self.addPoints(res,res)
				if bit == "1":
					res = self.addPoints(res,p)
			return res
		else:
			print "Must have a point on the curve as input"

	def addPoints(self,p1,p2):
		if self.checkPoint(p1) and self.checkPoint(p2):
			if p1 == self.id:
				return p2
			elif p2 == self.id:
				return p1
			elif (p1[0] == p2[0]) and (p1[1] == ((-p2[1])%self.p)):
				return self.id
			elif (p1[0] == p2[0]) and (p1[1] == p2[1]):
				return self.calcPoint(p1,p2,self.lambdaEqual(p1))
			else:
				return self.calcPoint(p1,p2,self.lambdaDiff(p1,p2))
		else:
			print "Both points must be on the curve."

	def inversePoint(self,p):
		if self.checkPoint(p) and (not p == self.id):
			return [p[0],((-p[1])%self.p)]
		elif p == self.id:
			return p
		else: 
			print "Must have point on curve to find inverse"

	def lambdaEqual(self,p1):
		l = ((3*p1[0]*p1[0])+self.a)*(self.modinv((2*p1[1]),self.p))%self.p
		return l

	def lambdaDiff(self,p1,p2):
		l = (p2[1]-p1[1])*(self.modinv(((p2[0]-p1[0])%self.p),self.p))%self.p
		return l

	def calcPoint(self,p1,p2,l):
		x = ((l*l) - p1[0] - p2[0])%self.p
		y = ((l*(p1[0]-x))-p1[1])%self.p
		return [x,y]

	def egcd(self,a, b):
	    if a == 0:
	        return (b, 0, 1)
	    else:
	        g, y, x = self.egcd(b % a, a)
	        return (g, x - (b // a) * y, y)

	def modinv(self, a, m):
	    g, x, y = self.egcd(a, m)
	    if g != 1:
	        raise Exception('modular inverse does not exist')
	    else:
	        return x % m

	
class ElGamal:

	rand = random.SystemRandom()

	def __init__(self,p=5):
		self.p = p

	def setGroup(self,ecc,params):
		self.ecc = ecc
		if self.ecc:
			self.ecc = ecc
			self.e = Ecc()
			self.e.setField(params[0])
			self.e.setCurve(params[1],params[2])
			self.numpoints = params[3]
		else:
			self.p = params[0]

	def keygen(self,base):
		if self.ecc:
			self.a = self.rand.randrange(0,self.numpoints)
			self.g = base
			self.h = self.e.multiplumOfPoint(base,self.a)
		else:
			self.a = self.rand.randrange(0,self.p)
			self.g = base
			self.h = self.square_and_multiply(self.g,self.a,self.p)

	def multiply(self,x,y):
		if self.ecc:
			return self.e.addPoints(x,y)
		else:
			return (x*y)%self.p

	def exponentiate(self,x,n):
		if self.ecc:
			if n < 0:
				x = self.e.inversePoint(x)
			return self.e.multiplumOfPoint(x,n)
		else:
			return self.square_and_multiply(x,n,self.p)

	def square_and_multiply(self, g, e, m):
		# g is the base, e is the exponent, m is the modulus.
		
		# We want to handle all integers e
		if e < 0:
			g = modinv(g)
			e = -e
		elif e == 0:
			return 1
			
		# Now assume that e > 0.

		# Get binary representation of e, and note that the first two
		# symbols are 0b when e is positive, so we cut them away.
		be = bin(e)[2:]
		
		a = 1
		for b in be:
			a = (a * a) % m
			if b == '1':
				a = (a * g) % m
		return a

	def encrypt(self,m):
		if self.ecc:
			t = self.rand.randrange(0,self.numpoints)
		else:
			t = self.rand.randrange(0,self.p)
		c1 = self.exponentiate(self.g,t)
		c2 = self.multiply(self.exponentiate(self.h,t),m)
		return [c1,c2]

	def decrypt(self,c):
		return self.multiply(c[1],self.exponentiate(c[0],(-self.a)))


print """

Usage: python -i elgamal.py
   or: ./elgamal.py (if you have set the executable flag)
	
The file provides two classes, Ecc and ElGamal. You can use Ecc by typing:

	e = Ecc()

It contains the following methods:

	setField(p): Sets the prime p of the field Z_p we want to work in.
	setCurve(a,b): Sets the values of a and b in the equation y^2 = x^3 + ax + b.
	checkPoint(p): Checks if a point p = [x,y] is a point on the curve over the given field, 
			and return True/False.
	multiplumOfPoint(p,n): Calculates np for a number n and a point p = [x,y], and returns 
			the result.
	addPoints(p,q): Returns the sum of the points p = [x1,y1] and q = [x2,y2]
	inversePoint(p): Returns the inverse of the point p on the curve. 
			               
Example
	e = Ecc()
	e.setField(11)
	e.setCurve(1,6)
	e.checkPoint([2,7])
	e.addPoints([2,7],[5,2])
	e.multiplumOfPoint([2,7],8)
	e.inversePoint([2,7])

You can use ElGamal by typing:
	cs = ElGamal()

It contains the following methods:
	setGroup(ecc,params): Sets the field. If ecc=True, generates an instance of the class
			Ecc to generate a elliptic curve to work over. Params the must be set as 
			params = [prime,a,b,numpoints] (prime sets field to work over, numpoint is number
			of points on the curve). If ecc=False, params should be params = [prime] to set
			the field Z_prime to work over.
	keygen(base): Sets g= base and generates a and h such that h = g^a.
	encrypt(m): Encrypts the message m (must be either point on curve or a number in the
			field.)
	decrypt(c): Decrypts the ciphertext c = [c1,c2].

Note: One setGroup is done, all operations will automatically happen over the selected group, 
	and you must ensure that inputs to keygen, encrypt and decrypt are valid for the given
	group.

Example
	cs = ElGamal()
	cs.setField(True,[11,1,6,13])
	cs.keygen([2,7])
	cs.encrypt([3,6])
	cs.decrypt([[10,2],[8,3]])

Note: If the encryption contains "Point at infinity", use cs.e.id as input for this point.

"""