# -*- coding: utf-8 -*-
#This file uses TABS

"""
This Module contains all the necessary Crypto-functions for MurrayB, including sha1-hashing, signing, signature-checking, encryption, decryption key-generation and keyring-handling
 """
import Crypto.Util.randpool
import Crypto.Hash.SHA
import Crypto.PublicKey.RSA
import Crypto.Cipher.AES
import thread
import cPickle
import service

__keyring_lock = thread.allocate_lock()
filename_identity = service.buildPath('identity.store')
filename_keyring = service.buildPath('keyring.store')
filename_murrayid = service.buildPath('murrayid.store')

def __save_identity_key(the_key):
	"""Saves the primary identity-key into a file"""
	try:
		fileToSave = file(filename_identity, 'w')
		cPickle.dump(the_key, fileToSave)
		fileToSave.close()
	except:
		print 'Something _very_ bad happened: could not write to "', filename_identity, '". Do you have the necessary rights and have ~3kB space left?'


def __load_identity_key():
	""" Loads the private Key if one exists. If not, it generates a new one. Returns the key."""
	
	try:
		fileRestoreFrom = file(filename_identity, 'r')
		identitykey_tmp = cPickle.load(fileRestoreFrom)
		fileRestoreFrom.close()
	except:
		identitykey_tmp = Crypto.PublicKey.RSA.generate(2048,randomnessgenerator.get_bytes)
		__save_identity_key(identitykey_tmp)
	return identitykey_tmp

def sha1hash( hashable_string ):
	"""returns the SHA1-hash of the input-string"""
	sha1hasher = Crypto.Hash.SHA.new()
	sha1hasher.update(hashable_string)
	tmp = sha1hasher.hexdigest()
	return tmp

def __save_keyring(keyring):
	"""Saves the keyring into a file"""
	try:
		fileToSave = file(filename_keyring, 'w')
		cPickle.dump(keyring, fileToSave)
		fileToSave.close()
	except:
		print 'Something _very_ bad happened: could not write to "', filename_keyring, '". Do you have the necessary rights and have space left?'


def __load_keyring():
	""" Loads the public Keyring if one exists. If not, it generates a new one. Returns the keyring as a dictionary."""

	try:
		fileRestoreFrom = file(filename_keyring, 'r')
		keyring_tmp = cPickle.load(fileRestoreFrom)
		fileRestoreFrom.close()
	except:
		tmpstring = sha1hash(cPickle.dumps(__identity_key.publickey()))
		keyring_tmp = { tmpstring : __identity_key.publickey() }
		__save_keyring(keyring_tmp)
	return keyring_tmp

def __load_murrayid():
	"""Loads the murrayid from a file. If the file does not exist, it creates a new murrayid."""
	try:
		fileRestoreFrom = file(filename_murrayid, 'r')
		murrayid_tmp = cPickle.load(fileRestoreFrom)
		fileRestoreFrom.close()
	except:
		murrayid_tmp = sha1hash(cPickle.dumps(__identity_key.publickey()))
		try:
			fileToSave = file(filename_murrayid, 'w')
			cPickle.dump(murrayid_tmp, fileToSave)
			fileToSave.close()
		except:
			print 'Something _very_ bad happened: could not write to "', filename_murrayid, '". Do you have the necessary rights and have ~50B space left?'
	return murrayid_tmp
		
def _aes_encrypt(text, password):
	"""Encrypts the text with AES. The password has to be a 256 bit long string."""
	aes_key = Crypto.Cipher.AES.new(password)
	paddingchar="X";
	if text[-1]=="X":
		 paddingchar="Y"
		 text = text + paddingchar
	else:
		if text[-1]=="Y":
			paddingchar="X"
			text = text + paddingchar
	if (len(text) % 16) <> 0:
		while (len(text) % 16) <> 0:
			text = text + paddingchar
	return aes_key.encrypt(text)

def _aes_decrypt(ciphertext, password):
	"""Decrypts the ciphertext with AES."""
	aes_key = Crypto.Cipher.AES.new(password)
	text = aes_key.decrypt(ciphertext)
	if text[-1] =="X":
		while text[-1] =="X":
			text= text[0:len(text)-1]
	else:
		if text[-1]=="Y":
			while text[-1] =="Y":
				text= text[0:len(text)-1]
	return text

def murrayid():
	"""returns the murrayid of the loaded identity"""
	return __murrayid

def known(murrayid_string):
	"""If the given Murrayid has a corresponding publickey, return true. Otherwise, return false"""
	try:
		__keyring_dictionary[murrayid_string]
		returnvalue = True
	except:
		returnvalue=False
	return returnvalue
	
def getPK():
	"""Returns our Public Key as a signed pickled string """
	thekey = cPickle.dumps(__identity_key.publickey())
	thesig = __identity_key.sign(thekey,'')
	return cPickle.dumps((thekey , thesig))

def storePK(murray_id , newpublickey):
	"""Stores a public-key inside the keyring. Returns True if everything went allright, False when errors occured"""
	if not known(murray_id):
		try:
			tmp=cPickle.loads(newpublickey)
			tmpkey=cPickle.loads(tmp[0])
			if (tmpkey.verify(tmp[0],tmp[1]))==1:
				__keyring_lock.acquire()
				__keyring_dictionary[murray_id]=tmpkey
				__save_keyring(__keyring_dictionary)
				__keyring_lock.release()
		except:
			return False
	return True
def sign(text):
	"""Signs the input text and returns the signature"""
	return cPickle.dumps(__identity_key.sign(sha1hash(text),''))

def verify_sig(text, signature,murrayid):
	"""verifies the signature. returns False if signature did not work out ok"""
	if known(murrayid):
		thekey = __keyring_dictionary[murrayid]
		if thekey.verify(sha1hash(text),cPickle.loads(signature)) == 1:
			return True
		else:
			return False
	else:
		return False

def encrypt(murray_id, plaintext):
	"""Signs and encrypts the plaintext"""
	signature = sign(plaintext)
	password = randomnessgenerator.get_bytes(32)
	randomnessgenerator.randomize();
	plainmessage=cPickle.dumps((plaintext,signature))
	ciphermessage=_aes_encrypt(plainmessage, password)
	publickey=__keyring_dictionary[murray_id]
	sessionkey = publickey.encrypt(password,"")
	endmessage = (sessionkey, ciphermessage)
	return cPickle.dumps(endmessage)

def decrypt(murray_id, ciphertext):
	"""Decrypts the plaintext and checks the signature. Returns None if the signature did not match"""
	try:
		endmessage=cPickle.loads(ciphertext)  #endmessage = (sessionkey, ciphermessage)
		sessionkey=__identity_key.decrypt(endmessage[0]) # sessionkey = password
		plaintextmessage=cPickle.loads(_aes_decrypt(endmessage[1],sessionkey)) # (plaintext,signature)
		if (verify_sig(plaintextmessage[0],plaintextmessage[1],murray_id)):
			return plaintextmessage[0]
		else:
			return None
	except:
		return None

randomnessgenerator = Crypto.Util.randpool.RandomPool(8192) #getting a cryptographic "good" randomnessgenerator
randomnessgenerator.randomize()
randomnessgenerator.stir()#stirring modifies the status of the generator to make it harder for an attacker to reconstruct it
__identity_key=__load_identity_key()
__keyring_dictionary=__load_keyring()
__murrayid =__load_murrayid()


"""TEST-Routine for Debugging 
If the module is executed directly, the following code will be executed
Basically, we use every function once.
"""
if __name__ == '__main__':
	errors = 0;
	print '==============================================='
	print 'Testing Crypto-module'
	print '==============================================='
	print ''
	print 'Testing randomnessgenerator'
	print '==============================================='
	print 'randomnessgenerator-size:', randomnessgenerator.bytes
	print 'randomnessgenerator-entropy:', randomnessgenerator.entropy
	print 'random bytes of the day:', randomnessgenerator.get_bytes(8)
	print 'randomnessgenerator-entropy:', randomnessgenerator.entropy
	print
	print 'Testing randomnessgenerator'
	print '==============================================='
	print 'sha1hasher will now hash the string "bacon"'
	print 'hashed bacon:',sha1hash('bacon')
	print 'should be   : 8abf15bef376e0e21f1f9e9c3d74483d5018f3d5'
	if sha1hash('bacon')<>'8abf15bef376e0e21f1f9e9c3d74483d5018f3d5':
		errors = errors +1;
	print
	print 'Testing RSA, AES + our identity'
	print '==============================================='
	print 'Our Murray-ID is:', murrayid()
	print 'Have we saved our Public Key in the keyring:', known(murrayid()), '(should be True)'
	if known(murrayid())<>True:
		errors = errors + 1;
	print 'I will now sign+encrypt \"bacon\" :', decrypt(murrayid(), encrypt(murrayid(),'bacon'))
	print 'This should of course be        : bacon'
	if decrypt(murrayid(), encrypt(murrayid(),'bacon'))<>'bacon':
		errors = errors + 1
	if errors <1:
		print 
		print '==============================================='
		print 'ALL TESTS COMPLETED SUCCESSFULLY'
		print '==============================================='
	else:
		print 
		print '==============================================='
		print 'Something went horribly wrong.'
		print '==============================================='
	raw_input()