''' Demonstrate Python AES library
	Author: Tim Pierson, Dartmouth CS55, Winter 2021
			based on https://github.com/kevin-w-du/BookCode/blob/master/Encryption/enc.py
	'''

from Crypto.Cipher import AES
from Crypto.Random import get_random_bytes
from Crypto.Util import Padding
block_size = 16 #AES uses 16 byte blocks

def encrypt(key, iv, message):
	print("Encrypting")
	#create new cipher using CBC 
	cipher = AES.new(key, AES.MODE_CBC,iv)
	#pad data to be a multiple of 16 bytes
	print("\tMessage length is",len(message))
	print("\tBlock size is",block_size)
	print("\tExpect padding to be",hex(block_size-len(message)%block_size))
	padded = Padding.pad(message,block_size)
	#create ciphertext
	ciphertext = cipher.encrypt(padded)
	return ciphertext

def decrypt(key, iv, ciphertext):
	#create new cipher using CBC 
	cipher = AES.new(key, AES.MODE_CBC,iv)
	plaintext = cipher.decrypt(ciphertext)
	return plaintext



if __name__ == '__main__':
	key = get_random_bytes(block_size)
	iv = get_random_bytes(block_size)
	message = b'We attack at dawn!'

	ciphertext = encrypt(key, iv, message)
	print("Ciphertext:")
	print(ciphertext)
	plaintext = decrypt(key, iv, ciphertext)
	print("\nDecrypted:")
	print(plaintext) #with padding
	print(plaintext.decode("utf-8")) #without padding