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

# Copyright (c) 2014, NORDUnet A/S.
# See LICENSE for licensing information.

import argparse
import urllib2
import urllib
import json
import base64
import sys
import struct
import hashlib
import itertools
from certtools import *
import zipfile

parser = argparse.ArgumentParser(description='')
parser.add_argument('baseurl', help="Base URL for CT server")
parser.add_argument('--store', default=None, metavar="dir", help='Store certificates in directory dir')
parser.add_argument('--start', default=0, metavar="n", type=int, help='Start at index n')
parser.add_argument('--single', default=None, metavar="n", type=int, help='Onlyfetch index n')
parser.add_argument('--verify', action='store_true', help='Verify STH')
args = parser.parse_args()

def extract_original_entry(entry):
    leaf_input =  base64.decodestring(entry["leaf_input"])
    (leaf_cert, timestamp, issuer_key_hash) = unpack_mtl(leaf_input)
    extra_data = base64.decodestring(entry["extra_data"])
    if issuer_key_hash != None:
        (precert, extra_data) = extract_precertificate(extra_data)
        leaf_cert = precert
    certchain = decode_certificate_chain(extra_data)
    return ([leaf_cert] + certchain, timestamp, issuer_key_hash)

def get_entries_wrapper(baseurl, start, end):
    fetched_entries = 0
    while start + fetched_entries < (end + 1):
        print "fetching from", start + fetched_entries
        entries = get_entries(baseurl, start + fetched_entries, end)["entries"]
        if len(entries) == 0:
            break
        for entry in entries:
            fetched_entries += 1
            yield entry

def print_layer(layer):
    for entry in layer:
        print base64.b16encode(entry)

sth = get_sth(args.baseurl)
tree_size = sth["tree_size"]
root_hash = base64.decodestring(sth["sha256_root_hash"])

print "tree size", tree_size
print "root hash", base64.b16encode(root_hash)

if args.single == None:
    entries = get_entries_wrapper(args.baseurl, args.start, tree_size - 1)
elif args.single > tree_size - 1:
    print "index", args.single, "too large, tree size:", tree_size
else:
    entries = get_entries_wrapper(args.baseurl, args.single, args.single)
    args.start = args.single

if args.verify:
    layer0 = [get_leaf_hash(base64.decodestring(entry["leaf_input"])) for entry in entries]

    tree = build_merkle_tree(layer0)

    calculated_root_hash = tree[-1][0]

    print "calculated root hash", base64.b16encode(calculated_root_hash)

    if calculated_root_hash != root_hash:
        print "fetched root hash and calculated root hash different, aborting"
        sys.exit(1)

else:
    currentfilename = None
    zf = None
    for entry, i in itertools.izip(entries, itertools.count(args.start)):
        try:
            (chain, timestamp, issuer_key_hash) = extract_original_entry(entry)
            if args.store:
                zipfilename = args.store + "/" + ("%04d.zip" % (i / 10000))
                if zipfilename != currentfilename:
                    if zf:
                        zf.close()
                    zf = zipfile.ZipFile(zipfilename, "w",
                                         compression=zipfile.ZIP_DEFLATED)
                    currentfilename = zipfilename
                s = ""
                s += "Timestamp: %s\n" % timestamp
                leaf_input = base64.decodestring(entry["leaf_input"])
                leaf_hash = get_leaf_hash(leaf_input)
                s += "Leafhash: %s\n" % base64.b16encode(leaf_hash)
                if issuer_key_hash:
                    s += "-----BEGIN PRECERTIFICATE-----\n"
                    s += base64.encodestring(chain[0]).rstrip() + "\n"
                    s += "-----END PRECERTIFICATE-----\n"
                    s += "\n"
                    chain = chain[1:]
                for cert in chain:
                    s += "-----BEGIN CERTIFICATE-----\n"
                    s += base64.encodestring(cert).rstrip() + "\n"
                    s += "-----END CERTIFICATE-----\n"
                    s += "\n"
                zf.writestr("%08d" % i, s)
        except AssertionError, e:
            print "error for cert", i, e
    if zf:
        zf.close()