# @file merkle_tree_4byte.py (2023-03-16T14:29Z)
# @author David Ireland <www.di-mgt.com.au/contact>
# @copyright 2023 DI Management Services Pty Ltd
# @license Apache-2.0

"""Generate a Merkle tree using weak 4-byte hash function."""

import hashlib

def H(val: str) -> str:
    """Weak hash function that returns a 4-byte hash.

    Args:
        val (str): hex-encoded input.

    Returns:
        str: Four-byte hash of input encoded in hex.
    """
    return hashlib.sha256(bytes.fromhex(val)).hexdigest()[:8]

def is_power_of_two(n: int) -> bool:
    return n > 0 and (n & (n -1)) == 0

def hash_double_arrays(hashes):
    """Compute hashes of set of double arrays.

    Args:
        hashes (str): Array of 2*n hex strings to be hashed pairwise.

    Returns:
        str: Array of n hex-encoded hash strings
    """
    n = len(hashes) // 2
    out = ['' for x in range(n)]
    for i in range(n):
        h = H(hashes[2 * i] + hashes[2 * i + 1])
        out[i] = h
    return out

def make_sk(seed, height, index):
    """Generate a private key value given a secret key and address."""
    adrs = "{:04x}".format(height) + "{:04x}".format(index)
    #print("adrs=", adrs)
    return H(seed + adrs)
  

# Generate keys
sk_seed = '900150983cd24fb0'
keys = [make_sk(sk_seed, 0, i) for i in range(8)]
print("keys =", keys)
print(f"size={len(keys)}")
assert(is_power_of_two(len(keys)))

# Hash the keys to leaf nodes
hashes = [H(k) for k in keys]
print("leaf nodes=", hashes)

# Compute the Merkle tree nodes
while (len(hashes) > 1):
    hashes = hash_double_arrays(hashes)
    print(hashes)

print(f"Root node is {hashes[0]}")