#!/usr/bin/python
# -*- coding: utf-8 -*-		

import time
import base64
import argparse
from pympler.asizeof import asizeof
from certtools import *


base_urls = ["https://plausible.ct.nordu.net/",
			"https://ct1.digicert-ct.com/log/", 
			"https://ct.izenpe.com/",
			"https://log.certly.io/", 
			"https://ct.googleapis.com/aviator/", 
			"https://ct.googleapis.com/pilot/",
			"https://ct.googleapis.com/rocketeer/",
			]

logkeys = {}
logkeys["https://plausible.ct.nordu.net/"] = get_public_key_from_file("../../plausible-logkey.pem")
logkeys["https://ct.googleapis.com/rocketeer/"] = get_public_key_from_file("../../rocketeer-logkey.pem")
logkeys["https://ct.googleapis.com/aviator/"] = get_public_key_from_file("../../aviator-logkey.pem")
logkeys["https://ct.googleapis.com/pilot/"] = get_public_key_from_file("../../pilot-logkey.pem")
logkeys["https://log.certly.io/"] = get_public_key_from_file("../../certly-logkey.pem")
logkeys["https://ct.izenpe.com/"] = get_public_key_from_file("../../izenpe-logkey.pem")
logkeys["https://ct1.digicert-ct.com/log/"] = get_public_key_from_file("../../digicert-logkey.pem")

parser = argparse.ArgumentParser(description="")
parser.add_argument('--audit', action='store_true', help="run lightweight auditor ensuring consistency in STH")
parser.add_argument('--build-sth', action='store_true', help="get all entries and construct STH")


def reduce_leafs_to_root(layer0):
    if len(layer0) == 0:
        return [[hashlib.sha256().digest()]]
    current_layer = layer0
    while len(current_layer) > 1:
        current_layer = next_merkle_layer(current_layer)
    return current_layer

# Get STH and verify signature
def fetch_all_sth():
	sths = {}
	for base_url in base_urls:
		try:
				sths[base_url] =  get_sth(base_url)
		except:
			print "Failed to retrieve STH from " + base_url
			continue

		try:
			check_sth_signature(base_url, sths[base_url], logkeys[base_url])
		except:
			print "Could not verify signature from " + base_url + "!!!"
			continue
	return sths


def verify_consistency(old, new):
	for url in old:
		# try:
			if old[url]["tree_size"]!= new[url]["tree_size"]:
				consistency_proof = get_consistency_proof(url, old[url]["tree_size"], new[url]["tree_size"] )
				decoded_consistency_proof = []
				for item in consistency_proof:
					decoded_consistency_proof.append(base64.b64decode(item))
				res = verify_consistency_proof(decoded_consistency_proof, old[url]["tree_size"], new[url]["tree_size"], old[url]["sha256_root_hash"])
				
				if old[url]["sha256_root_hash"] != str(base64.b64encode(res[0])):
					print "Verification of old hash failed!!!"
					print old[url]["sha256_root_hash"], str(base64.b64encode(res[0]))
				elif new[url]["sha256_root_hash"] != str(base64.b64encode(res[1])):
					print "Verification of new hash failed!!!"
					print new[url]["sha256_root_hash"], str(base64.b64encode(res[1]))
				else:
					print time.strftime("%H:%M:%S", time.gmtime()) + " New STH from " + url + ", timestamp: " + str(new[url]["timestamp"]) + ", size: " + str(new[url]["tree_size"]) + "...OK."

		# except:
		# 	print "ERROR: Could not verify consistency for " + url


def fetch_and_build_tree(old_sth, base_url):
	print "Getting all entries from " + base_url

	sth = old_sth[base_url]

	entries = []
	leafs = []

	while len(leafs) < sth["tree_size"]:
		pre_size = len(leafs)
		entries = get_entries(base_url, len(leafs), sth["tree_size"])["entries"]
		
		for item in entries:
			leafs.append(get_leaf_hash(base64.b64decode(item["leaf_input"])))
		print "Got entries " + str(pre_size) + " to " + str(len(leafs)) + " (total leaf size: " + str(asizeof(leafs)/1024/1024) + "MB)"


	# tree = build_merkle_tree(leafs)
	root = base64.b64encode(reduce_leafs_to_root(leafs)[0])
	# print "Tree size: " + str(asizeof(tree)/1024/1024) + "MB"
	# print len(leafs)

	if root == sth["sha256_root_hash"]:
		print "Verifying root hashes...OK."
	else:
		print "ERROR: Failed to verify root hashes!"
		print "STH root: " + sth["sha256_root_hash"]
		print "Tree root: " + root


def main(args):
	print "Started " + time.strftime("%H:%M:%S", time.gmtime())
	old_sth = fetch_all_sth()

	if args.build_sth:
		fetch_and_build_tree(old_sth, base_urls[2])

	if args.audit:
		print "Running auditor for " +str(len(base_urls)) + " logs..."

		while True:
			time.sleep(1*60-4)
			# print time.strftime("%H:%M:%S", time.gmtime()) + " checking for updates..."
			new_sth = fetch_all_sth()
			verify_consistency(old_sth, new_sth)
			old_sth = new_sth





if __name__ == '__main__':
    main(parser.parse_args())