# -*- coding: utf-8 -*-

import string
import pickle
from Cryptodome.Cipher import AES
from Cryptodome.Random import get_random_bytes

plain = b'Homework 4 - RT!'
pwchar = string.ascii_letters + string.digits + "!#"
pwlen = 10
rainbow_const = []

def ba2pass(ba, column):
	if len(ba) != 16:
		raise ValueError('Byte array length must be 16.')
	xored = [a^b for (a,b) in zip(ba,rainbow_const[column])]
	return bytes([ord(pwchar[xored[i] % 64]) for i in range(pwlen)])

def gen_chain(start, t):
	pw_val = start
	for column in range(t):
		pw_val = pw_val + b'0000000000000000'[:16-len(pw_val)]
		aes = AES.new(pw_val, AES.MODE_ECB)
		ct = aes.encrypt(plain)
		pw_val = ba2pass(ct, column)
	return (start, ct)

def rainbowtable(t = 2**12, m = 2**14):
	print('GENERATING Rainbow table')
	# initialize some rainbow constants
	aes = AES.new(b'0000000000000000', AES.MODE_ECB)
	for i in range(t):
		hs = '{0:x}'.format(i)
		if len(hs) % 2 == 1:
			hs = '0' + hs
		x = bytes.fromhex(hs)
		rainbow_const.append(aes.encrypt(x + b'0000000000000000'[:16-len(x)]))
	# generate rainbow table
	RT = {}
	for i in range(m):
		if (i%100 == 0):
			print(i)
		seed = ba2pass(get_random_bytes(16), 0)
		(start,end) = gen_chain(seed, t)
		RT[end] = start
	return RT

def writeRT(RT, name):
	with open(name, "wb") as f:
		pickle.dump(RT, f)

def readRT(name):
	with open(name, "rb") as f:
		RT = pickle.load(f)
	return RT


if __name__ == "__main__":
	t = 2**10
	m = 2**16
	RT = rainbowtable(t, m)
	writeRT(RT, "RT.dat")
	RT = {}
	RT = readRT("RT.dat")
	print('-'*40)