# --------------------------------------------------
#
#   Copyright (c) 2003-2005 by rattle.
#             .aware security research
#
#   http://www.awarenetwork.org/
#
# --------------------------------------------------
#
#   This is pyECC, an ECC implementation in Python   
#
#   pyECC is free software; you can redistribute it
#   and/or modify it under the terms of the GNU
#   General Public License as published by the Free
#   Software Foundation; either version 2 of the
#   License, or (at your option) any later version.
#
#   pyECC is distributed in the hope that it will be
#   useful, but WITHOUT ANY WARRANTY; without even
#   the implied warranty of MERCHANTABILITY or 
#   FITNESS FOR A PARTICULAR PURPOSE. See the GNU 
#   General Public License for more details.
#
#   You should have received a copy of the GNU
#   General Public License along with Plague; if not,
#   write to the Free Software Foundation, Inc.,
#
#   59 Temple Place,
#   Suite 330, Boston,
#   MA 02111-1307 USA
#
# --------------------------------------------------

from math import *
from aes import *
from base64 import encodestring as b64
from random import random


# This is a randomization routie that uses
# random() from the module random, but returns
# a long int instead of a float.
def rand(maximum): 
	fl = random()
	ll = long(fl * (10**17)) # thats the maximum precision
	ll *= maximum
	ll /= (10**17)
	return ll


# Generate a "good" prime which satisfies the 
# requirement that 4 divides p+1. This is required
# for calculating the square root inside a finite
# field modulo p.
def primeD4(bytes):
	p = 1
	while ((p+1)%4): p = prime(bytes,0)
	return p


# This procedure generates a large prime number
def prime(bits, init=0L):
	i = init
	if i: i+= long(log(i)/2)
	else: i = rand(2**bits)
	if not i%2: i+=1

	# p is the number of rounds we will let 
	# the Miller-Rabin algorithm pass for 
	# verifying that our number is prime.
	# A maximum of 40 checks makes sense to
	# not slow down the algorithm unnecessariely.
	p = int(ceil(sqrt(bits)))*2
	if (p > 40): p = 40

	f = False 
	while not f:
		while not Fermat(i): i += 2	# first check with fermat
		f = MRabin(i,p)			# verify with Miller-Rabin
	return i

def Fermat(n):
	return n==2 or F(n,2)**(n-1)==1


# The Miller-Rabin algorithm is a famous little 
# trick to verify that a pretty large number is
# prime. Basically, we ask a certain number of
# random "wintnesses" (integer numbers) whether
# they can verify that our number is not prime.
# If we cannot find any such witness, we assume
# that number is prime. If you want to know why
# this is extremely reliable, read up on it.
def MRabin(number,attempts):
	rndNum = 0L
	retVal = False
	i = 0
	retVal = True;
	for i in xrange(attempts):
		rndNum = rand(number-2)		
		if (rndNum < 2): rndNum = rndNum + 2
		if (Witness(rndNum, number)):
			retVal = False
			break
	return retVal
def Witness(witness,number):
	f = 1; x = 0;
	retVal = False;
	n = number - 1
	i = long(ceil(log(n)/log(2))+1)
	while (i >= 0):
		x = f
		f = x * x % number
		t = 2 ** i
		if (f == 1) and (x != 1) and (x != n):
			retVal = True
			break
		if (n >> i) & 1:
			f = f * witness % number;
		i -= 1
	if (retVal):
		return True
	else:
		if (f != 1): return True
		else: return False
		


# The classes lrange and i_lrange implement a range() object
# for realy large, long integers. Even xrange is still based
# on int objects, not long objects.

class i_lrange:
	def __init__(self,m,n,s):
		self.n = long(n)
		self.m = long(m)
		self.s = long(s)
	def __iter__(self):
		return n
	def next(self):
		n = self.n
		self.n += self.s
		if self.n>self.m: 
			raise StopIteration
		return n
class lrange:
	def __init__(self,m,n=0,stp=1):
		self.max = long(max(m,n)) + 1L
		self.min = long(min(n,m))
		self.stp = long(stp)
	def __iter__(self):
		return i_lrange(self.max,self.min,self.stp)
	def __contains__( self, item):
		if (item <= self.max) and (item >= self.min):
			return (item-self.min) % self.stp == 0
		return False
	def __len__(self):
		return int((self.max - self.min) / self.stp)




# The class F implements elements of a finite field over a
# given number m. As you should know, m has to be prime
# in order for F_m to be a field.
# However, this class does not check whether m is prime so
# you can also use this cass for arbitrary finite sets.
# I implemented all aritmethic operations modulo m, and 
# there is a fast __pow__ variant as well. If you divide
# by an instance of this class, you actually multiply
# with the multiplicative inverse, which is calculated
# by using euler's extended algorithm. You can also use
# the unary operator ~ to get the multiplicative inverse.

class F:
	def __init__(self, m, a):
		try:
			self.mod = abs(long(m))
			self.num = abs(long(a))
		except:
			raise ValueError("F needs numeric parameters")
		else:
			self.__c()

	def __r(self,a):
		return F(self.mod,a) 
	def __c(self):
		self.num %= self.mod

	def __long__(self): 
		return self.num
	def __float__(self):
		return float(self.num)

	def __cmp__(self,other):
		x = F(self.mod, long(other))
		if self.num < x.num:
			return -1
		elif self.num == x.num:
			return 0
		else:
			return 1

	def __pow__(self,b):
		if (b < 0):
			raise ValueError("Use mul_inv() before ModExp")
		elif (b == 0):  
			return self.__r(1)
		elif (b < 1):
			if (b == 0.5):
				x = divmod(self.mod+1,4)
				x = x[0] - int(bool(x[1]))
	 			return self.__pow__(x)
			else:
				raise ValueError("Cannot return this root.")

		b = long(b)
		d = self.__r(1)
		p = long(ceil(log(b)/log(2))+1)
		for i in lrange(1,p):
			d = d * d
			if (b>>(p-i))&1: d *= self
		return d


	def __neg__(self):
		return (self.mod - self.num)

	def __invert__(self):
		m = self.mod
		a = self.num
		q=0L; r=0L;
		x = [0L,0L,0L]	
		y = [0L,0L,0L]
		x[2] = 1; x[1] = 0
		y[2] = 0; y[1] = 1
		while (m>0):
			q=a/m
			r=a-q*m
			x[0]=x[2]-q*x[1];
			y[0]=y[2]-q*y[1]
			a,m=m,r
			x[2]=x[1];x[1]=x[0];
			y[2]=y[1];y[1]=y[0];
		if x[2]<0: x[2] = self.mod+x[2]
		return x[2]

	def __add__(self,other):
		return self.__r(self.num + long(other))
	def __sub__(self,other):
		inverse = F(self.mod, other)
		inverse = -inverse
		return self + inverse
	def __mul__(self,other):
		return self.__r(self.num * long(other))
	def __div__(self,other):
		inverse = F(self.mod, other)
		inverse = ~inverse
		return self * inverse

	def __radd__(self,other):
		return self.__r(self.num + long(other))
	def __rmul__(self,other):
		return self.__r(self.num * long(other))

	def __rdiv__(self,other):
		return F(self.mod, other).__div__(self)
	def __rsub__(self,other):
		return F(self.mod, other).__sub__(self)
	
	def __repr__(self):
		return '%d' % self.num

	def legendre(self):
		l = self.__pow__((self.mod-1)/2)
		if l and l!=1: return -1
		return l

	def __hex__(self):
		return hex(self.num)
	
	def __hash__(self):
		return hash(self.num)




# This class is the iterator for elliptic curves over finite
# prime fields. It returns the points which satisfy the curve
# equation.

class citer:
	def __init__(self, c, x):
		self.c = c
		self.x = x
		self.t = False
	def __iter__(self):
		return self
	def __ck(self):
		self.x += 1
		if (self.x == 0): 
			self.t = True
	def next(self):
		if self.t:
			raise StopIteration
		while self.c.y2(self.x).legendre() != 1:
			self.__ck()
			if self.t: raise StopIteration

		p = self.c.point(self.x)
		self.__ck()
		return p


# This is the class for elliptic curves. 

class elliptic_curve:
	def __init__(self, a, b):
		self.a = a
		self.b = b

	def y2(self,x):			# calculate y to the 
		r  = self.b		# power of 2 at x
		r += self.a * x
		r += x ** 3
		return r

	def y(self,x):			# return y (assumes that a 
		z = self.y2(x)		# square root of y square
		return (z**0.5)		# esists).

	def point(self,x):			# return the point 
		return P(x,self.y(x),self)	# at x.

	def __eq__(self,other):
		return (self.a,self.b) == (other.a,other.b)
	def __ne__(self,other):
		return (self.a,self.b) != (other.a,other.b)
	def __len__(self):
		return int(self.order())
	def __contains__(self,p):
		return p[1]**2 == self.y2(p[0])
	def __repr__(self):
		args = (long(self.a), long(self.b))
		return "<elliptic curve: y\xFD=x\xFC+%dx\xFD+%d>" % args


# The following class inherits from elliptic_curve and 
# implements curves over finite prime fields. This is
# the curve we will be using for cryptographic purposes.

class finite_field_curve(elliptic_curve):
	def __init__(self, m, a=0, b=0):
		while 0 == (4*(a**3) + 27*(b**2)):
			if not a and not b:
				a = rand(m)
				b = rand(m)
			else:
				raise ValueError("Invalid Elliptic Curve: 4a\xFC + 27b\xFD = 0")
		elliptic_curve.__init__(self,F(m,a),F(m,b))
		self.m = m
	def __iter__(self):
		return citer(self,F(self.m,0))
	def random(self):
		n = min(self.m+1 - 2*int(floor(self.m**0.5)), self.m+1 + 2*int(floor(self.m**0.5)))
		n = rand(self.m-n)+n
		while self.y2(n).legendre() != 1:
			n += 1
		return self.point(n)




# The P class implements a point on an elliptic curve. Thanks to
# the flexibility of python's data types, this class does not
# necessariely implement points over a finite field - the point
# merely satisfies the curve equation. 
#
# Point addition and scalar multiplication have been implemented
# as overloaded operators. 

class P:
	def __init__(self,x,y,curve):
		self.curve = curve
		
		if x is None and y is None:
			self.x = self.y = None
		else:
			self.x = x
			self.y = y
			
			if self.curve.y2(x) != (self.y ** 2):
				raise ValueError("(%d,%d) doesnt seem to be in curve." % (long(x),long(y)))


	# Two points are equal when they have the 
	# same coordinates. __coerce__ and __getitem__
	# have been added for fun and comfort. 

	def __eq__(self,other):	
		return (tuple(self)==tuple(other))
	def __ne__(self,other):
		return (tuple(self)!=tuple(other))

	def __coerce__(self,other):
		try: 
			x,y = tuple(other)
			return (self, P(x,y,self.curve))
		except:
			return None

	def __getitem__(self,i):
		if i and i!=1: raise IndexError 
		else: return long(i*self.y) or long(self.x)


	# Returns true if this is the point at inifinity.

	def inf(self):
		return self.x is None and self.y is None	

	
	# scalar multiplication is commutative just because
	# we define it that way.	
	def __rmul__(self,other):
		return self.__mul__(other)

	# Scalar multiplication has been implemented using 
	# the Square & Multiply (or rather Double & Add) 
	# method. If you haven't seen this a billion times
	# before, check wikipedia.
	#
	# Note: This is the actual important part of elliptic
	# curve cryptography. The discrete logarithm problem
	# occurs when trying to find z so that z*P=Q for known
	# points P and Q. Calculating the scalar multiple
	# of z*P is fairly easy and fast, but inverting this
	# process is ... let's say "less" easy. And that's
	# where ECDLP key exchange kicks in.

	def __mul__(self, a):
		try: 
			a = abs(long(a))
		except:	
			raise ValueError("Cannot multiply with another point")
		
		pt = P(None,None,self.curve)

		if a:
			pw = long(ceil(log(a)/log(2))+1)
			for i in lrange(1,pw):
				pt = pt + pt
				if (a >> (pw-i)) & 1: 
					pt = pt + self
		return pt


	# The additive inverse of a point is (x,-y). Yea.
	# Well that's just the way it is.
	def __neg__(self):
		if self.inf(): return self
		else: return P(self.x, -self.y, self.curve)

	# subtracting means adding the additive inverse. 
	def __sub__(self,other):
		return self.__add__(-other)
	
	# The addition is the most interesting part about the
	# whole elliptic curve cryptography study. Less abstractly
	# spoken, we draw a line through the two points that are
	# being added. If this line then intersects the curve at
	# a point Q, we calculate -Q and use this point as the 
	# result of our addition. If the line doesn't intersect 
	# the curve in another point, we say that it instersects
	# it in a point "at inifinity" which will be the neutral
	# element of the addition declared this way. And well.
	# The points on the curve with this addition are a commutative
	# group. Trust me, they are!

	def __add__(self, other):
		m = x1 = y1 = None
		if self.inf():
			return other
		elif other.inf():
			return self
		elif self.x == other.x:
			if (self.y+other.y) == 0:
				x1,y1 = None,None
			elif other.y != 0 and self.y != 0:
				m = (3*(self.x**2)+self.curve.a) / (2*self.y)
		else:
			m = (other.y - self.y) / (other.x - self.x)
		if m is not None:
			x1 = (m**2) - self.x - other.x
			y1 = (-m * (x1 - self.x)) - self.y

		return P(x1,y1,self.curve)

	def __repr__(self):
		if self.inf(): return "(*,*)"
		else: return repr((self.x, self.y))

	def __str__(self):
		return str(self.x)+str(self.y)
	


# This function returns the base point, which will be
# used for the key exchange procedure.

def random_point(b):
	m = primeD4(b)
	return finite_field_curve(m).random()

	
# The following class implements the key itself. As you should
# know, we use the ECDLP (elliptic curve discrete logarithm problem)
# for the key exchange. The actual encryption should then be 
# performed using AES (or a similarly strong and sufficiently fast
# algorithm). Below you can see a little sample session which might
# straighten up how it is supposed to work:
#
# >>> from ecc import *
# >>> p = random_point(164)
# >>> e,d=key(p),key(p)
# >>> e.exchange(d.public)
# "\x9f?\xb6<\x1f\x95\xbel\n\xf3\x8bi\xfcs{g\xcb\x7f\x80-\x02%\xed\x93.V\x88'eUY\x1aq$\t\xa7\xf3\xc9@\xc3\xed\x02"
# >>> e.exchange(d.public) == d.exchange(e.public)
# True
# >>> e
# HDKLPKCukSzTa/YVaRFjKaVzVs6wnQWoozEGP6c0AQxdgzmXGcHBhw4Fq7/VJZHiteBhhVzrshfG
# 5FsLMdPJ4b8hc2ETJMjpEdcMmaCw0EuRXg==
# 
# >>> d
# 2Xr/kx6yW85b/eVY6xNSKbbgoaQB7FLty5EWRXgEHpEMQCc2hzDcIsUCm0Fq9OEg/JAqCy0UmeYP
# FtuOtvgGKHFDlpnJlYCmsvspHY/4GoMAJP4m
# 
# >>> aes(d.exchange(e.public)[:32]).decipher(
# ... aes(e.exchange(d.public)[:32]).encipher("Hello World!"))
# 'Hello World!\x00\x00\x00\x00'
# >>>


class key:
	def __init__(self,p):
		self.private = rand(2**(16*self.bytelen(p.x)))
		self.public  = self.private * p

	def exchange(self, public):
		return self.point_hex(self.private * public)

	def point_hex(self,q):
		return self.hexlify(q.x)+self.hexlify(q.y)

	def bytelen(self,num): 
		i = long(num)
		blen = 0
		while (i != 0):
			blen += 1 
			i >>= 8 
		return blen

	def hexlify(self,num):
		sval = ""
		l = self.bytelen(num)
		i = long(num)
		for j in range(l):
			sval += chr(i&0xFF)
			i >>= 8
		return sval
	
	def __repr__(self):
		return b64(self.point_hex(self.public)+self.hexlify(self.private))