"""
**Encrypted socket class**
"""
#====================================================================
# cryptsock python module
#
# Copyright (c) 2001, Bryan Mongeau <bryan@eevolved.com>
# All rights reserved.
#
# This code is hereby placed in the public domain.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1- Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
#
# 2- Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
#
# 3- The names of its contributors may not be used to endorse or promote
#    products derived from this software without specific prior written
#    permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE REGENTS OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
#  NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
#====================================================================

import socket
from zlib import compress, decompress
from entropy.entropy import CSPRNG
from ecc.ecc import ecc
from aes.aes import aes
from cPickle import *

class cryptsock:
	"""
	This class provides a transparent layer of security
	for network transmissions when properly used. In this
	and other respects, it draws much from the goals of
	SSL and TLS. However, cryptsock's underlying cryptographic
	primitives are more efficient and consequently, impose
	less overhead in terms of memory, cycles and module size.

	This class basically wraps the traditional socket, providing
	a few new methods to be used. The raw socket is accessible as
	the *sock* attribute of the class. Zlib compression
	is used on all bulk data that passes through the socket.

	Uses 128 Bit AES in ECB mode, Elliptic curve Diffie-Hellman
	and Nyberg-Rueppel Signature / Verification. A record
	protocol based on TLS (rfc 2246) and cPickle is	employed
	to manage message integrity.

	**Sample usage: Echo program**::

		Server:
		-------

		from cryptsock.cryptsock import cryptsock
		s=cryptsock(1)
		s.bind("",12345)
		s.listen(1)
		print 'Listening...'
		c=s.accept()
		print 'Connect from ', c.addr
		data = c.recv()
		c.send(data)
		c.close()

		Client:
		-------

		from cryptsock.cryptsock import cryptsock
		s=cryptsock()
		s.connect('127.0.0.1',12345)
		s.send('Hello World!')
		data = s.recv()
		s.close


	As you can see, cryptsock's API is nearly identical to
	the standard python socket.

	This module requires the AES, ECC and CSPRNG classes
	from cryptkit in order to function properly.
	"""
	def __init__(self, server=0, sock=None, family=socket.AF_INET, type=socket.SOCK_STREAM):
		"""
		Constructor.
		Instantiates a CSPRNG and keys the ECC.
		Also instantiates an AES cipher for use
		afterwards.
		*sock*, if provided, constitues the socket to use,
		otherwise creates a new one.

		*server*, if true, will bypass the creation of
		crypto primitives for this instance, creating them instead
		for the socket instances it spawns in its accept() loop.

		*family* and *type* are the standard
		socket constructor arguments.
		"""
		if not server:
			self.oracle = CSPRNG()
			self.ecc = ecc(self.oracle.getInt())
			self.aes = aes(None,"ECB")

		if sock:
			self.sock=sock
		else:
			self.sock = socket.socket(family,type)
			self.sock.setsockopt(socket.SOL_SOCKET,socket.SO_REUSEADDR,1)

	def bind(self,host,port):
		"Just a dumb wrapper for the underlying socket."
		self.sock.bind((host,port))
		self.addr=host

	def listen(self,num):
		"Just a dumb wrapper for the underlying socket."
		self.sock.listen(num)

	def close(self):
		"Just a dumb wrapper for the underlying socket."
		self.sock.close()

	def reKey(self,key):
		" Re-Key the symmetric cipher. "
		self.aes.setKey(key)

	def accept(self):
		"""
		Here is where the server-side of the key exchange
		takes place.
		Spawns a brand new instance of cryptsock.
		"""
		conn, addr = self.sock.accept()
		ret = cryptsock(0,conn)
		ret.pk = loads(conn.recv(ret.ecc.pubKeySize))
		conn.send(dumps(ret.ecc.publicKey(),1))
		ret.addr = addr
		ret.aes.setKey(ret.ecc.recv(ret.pk)[:16])
		return ret

	def connect(self, host, port):
		"""
		Here is where the client side of the key exchange takes
		place.
		"""
		self.sock.connect((host,port))
		self.sock.send(dumps(self.ecc.publicKey(),1))
		self.pk = loads(self.sock.recv(self.ecc.pubKeySize))
		self.aes.setKey(self.ecc.recv(self.pk)[:16])

	def send(self,msg):
		"""
		Sends the arbitrary *msg* over the encrypted
		connection. First it places the message in a parcel
		containing a signature. Then it pickles, compresses,
		encrypts and sends that parcel over the socket.
		"""
		parcel=(msg,self.ecc.sign(dumps(msg,1)))
		self.sock.send(self.aes.encrypt(compress(dumps(parcel,1))))

	def recv(self,maxBytes=1024):
		"""
		Receives up to *maxBytes* bytes of data from the
		socket. Defaults to 1024. Method returns the unencrypted,
		decompressed, unpickled version of the message.
		"""
		data=self.sock.recv(maxBytes)
		if data:
			parcel=loads(decompress(self.aes.decrypt(data)))
			if self.ecc.verify(dumps(parcel[0],1),self.pk,parcel[1]):
				return parcel[0]
			else:
				self.sock.close()
				raise RuntimeError ("Verification has failed. Socket has been closed.")
		else:
			self.sock.close()
			raise RuntimeError ("Socket has been remotely closed.")