#!/usr/bin/env python

import argparse
from sys import exit
from Crypto.Hash import SHA256
from random import randint
from time import time
from pickle import dump
from binascii import b2a_hex

alpha = list("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789./")

# update() - produce next password in sequence
def update(pw, alpha):
	if len(pw) == 0:
		return ""
	x = alpha.index(pw[0])
	x = (x + 1) % len(alpha)
	if (x == 0):
		pw = alpha[x] + update(pw[1:], alpha)
	else:
		pw = alpha[x] + pw[1:]
	return pw

# gen_Rmask() - generate masks for return functions
#   R(i, hash) = R0(hash ^ Rmask[i])
def gen_Rmask(num_masks):
	Rmask = []
	sha = SHA256.new()
	for i in range(num_masks):
		sha.update(str(i))
		Rmask.append(sha.hexdigest())
	return Rmask
			    
# R() - return function
def R(hash, Rmask, alpha, length):

	alpha_mask = pow(2, 6) - 1
	pw = ""

	# xor the Rmask with the hash
	bhash = bytearray.fromhex(hash)
	bRmask = bytearray.fromhex(Rmask)
	for i in range(len(bhash)):
		bhash[i] = bhash[i] ^ bRmask[i]
	hash = b2a_hex(bhash)

	# generate the password
	for i in range(length):
		index = alpha_mask & ord(hash[(2 * i):(2 * (i + 1))].decode('hex'))
		pw = pw + alpha[index]

	return pw

def gen_string(pw, alpha, Rmask, prefix_set, prefix_len, max_string_len):

	pw_length = len(pw)
	string_length = 1
	start = pw

	sha = SHA256.new()
	sha.update(pw)
	digest = sha.hexdigest()

	while digest[:prefix_len] not in prefix_set:
		pw = R(digest, Rmask, alpha, pw_length)
		string_length = string_length + 1
		if string_length > max_string_len:
			return start, "", 0
		sha = SHA256.new()
		sha.update(pw)
		digest = sha.hexdigest()

	return start, digest, string_length

def random_password(length):
	null_Rmask = "00"*length
	x = hex(randint(0, (1 << length*8) - 1))[2:]
	while len(x) < 2*length:
		x = "0" + x
	return R(x, null_Rmask, alpha, length)


if __name__ == '__main__':


						      
	# password length
	pw_length = 5

	# prefix set for distinguished endpoint test
	prefix_set = set(["000", "001", "002", "003"])
	prefix_len = 3

	# TMTO parameters
	t = pow(2, 10)  # string length
	r = pow(2, 10)  # number of tables
	m = pow(2, 10)  # number of strings per table

	parser = argparse.ArgumentParser()
	parser.add_argument('--start-block', nargs='?', type=int, dest='start_block', default=0)
	parser.add_argument('--end-block', nargs='?', type=int, dest='end_block', default=r-1)
	args = parser.parse_args()
	
	if args.start_block > args.end_block or args.start_block > r or args.end_block > r:
		print "[-] Invalid start or end blocks on command line"
		exit(1)
	
	# catalog file base name
	outfile_base = "catalog_30bit_v2"

	Rmask = gen_Rmask(r)

	start_time = time()

	print "[*] Beginning table build..."

	for i in range(args.start_block, args.end_block+1):
		count = 0
		catalog = {}
		print "[*] Beginning block build with return function "+ str(i)
		block_start_time = time()
		while count < m:
			if (count % 128 == 0):
				print "[*] Return function " + str(i) + ", string count " + str(count)
			pw = random_password(5)
			pw, digest, string_len = gen_string(pw, alpha, Rmask[i], prefix_set, 
							    prefix_len, 3*t)
			if string_len > 0 and digest not in catalog.keys():
				count = count + 1
				catalog[digest] = pw
		
		elapsed_time = time() - block_start_time

		print "[*] Block " + str(i) + " build time: " + str(elapsed_time) + " seconds"

		outfile = outfile_base + ".block" + str(i)

		print "[*] Writing block " + str(i) + " to file " + outfile

		fp = open(outfile, "w")
		dump(catalog, fp)
		fp.close()

	elapsed_time = time() - start_time

	print "[*] Elapsed time: " + str(elapsed_time/3600) + " hours"
	print "[*] Average time per string: " + str(elapsed_time/(m*r)) + " seconds"