from Crypto.Random import get_random_bytes
from Crypto.Cipher import AES
import base64
import sys

# Encrypt and decrypt using pycrypto AES:
# - AES MODE_CBC symmetric cypher
# - 256 bit key, urlsafe_base64 encoded
# - random 128 bit nonce
# - PKCS7 padding
# - cyphertext is returned urlsafe_base64 encoded
# - fast and compact

def gen_key():
	key = get_random_bytes(256 / 8)
	key_b64 = base64.urlsafe_b64encode(key)
	return key_b64

def pad_pkcs7(s):
	length = 16 - len(s) % 16
	p = s + chr(length) * length
	return p

def unpad_pkcs7(p):
	s = p[:-ord(p[-1])]
	return s

def encrypt(plaintext, key_b64):
	key = base64.urlsafe_b64decode(key_b64.upper())
	assert len(key) == 32, "key is not a 256-bit base16 encoded string"
	nonce = get_random_bytes(AES.block_size)
	cipher = AES.new(key, AES.MODE_CBC, nonce)
	encrypted = cipher.encrypt(pad_pkcs7(plaintext))
	ciphertext = nonce + encrypted
	ciphertext_b64 = base64.urlsafe_b64encode(ciphertext)
	return ciphertext_b64

def decrypt(ciphertext_b64, key_b64):
	ciphertext = base64.urlsafe_b64decode(ciphertext_b64)
	key = base64.urlsafe_b64decode(key_b64.upper())
	assert len(key) == 32, "key is not a 256-bit base16 encoded string"
	nonce = ciphertext[0:AES.block_size]
	encrypted = ciphertext[AES.block_size:]
	cipher = AES.new(key, AES.MODE_CBC, nonce)
	plain = unpad_pkcs7(cipher.decrypt(encrypted))
	return plain

def test(n=2):
	crypto = sys.modules[__name__]
	key = crypto.gen_key()
	print "key", key

	text = "hello \0 world, this is a little crypto test\0"
	print "text, len", repr(text), len(text)
	for i in range(0, int(n)):
		enc = crypto.encrypt(text, key)
		print "encrypted", enc
		dec = crypto.decrypt(enc, key)
		assert dec == text, "crypto.test"
		print "decrypted OK"

if __name__ == "__main__":
	op = sys.argv[1]
	print globals()[op](*sys.argv[2:])