# @file wots_PKgen.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

"""Compute WOTS secret key then derive public key."""

from spx_sha256 import F, PRF

def setADRS(adrs_base, height, index, iset):
    """Crude way to set ADRS string."""
    t = 64
    treeindex = iset * t + index
    return adrs_base + format(height, f'08x') + format(treeindex, f'08x')

w = 16
# Working in the bottom HT subtree...
# ADRS excluding last two 4-byte words
adrs_base = '0028daecdc86eb87610000000000'
PKseed = 'B505D7CFAD1B497499323C8686325E47'
SKseed = '7c9935a0b07694aa0c6d10e4db6b1add'

print("Generate WOTS+ private key for i = 0")
# NB This is not the same as the sk used to generate the signature!
# sk = PRF(SK.seed, ADRS)
adrs_c = setADRS(adrs_base, 0, 0, 0)
print(f"ADRS={adrs_c}")
sk = PRF(SKseed, adrs_c)
print(f"sk={sk}")
#print(f"OK=c04623124dfcdcb1de0ad8cfc68ebf73")

# Compute F^w(sk)
# NB We only do this (w-1) times!
#adrs_base = '0028daecdc86eb87610000000000'
x = sk
for i in range(w-1):
    adrs_c = setADRS(adrs_base, 0, i, 0)
    print(f"i={i} ADRS={adrs_c}")
    print(f"in={x}")
    x = F(PKseed, adrs_c, x)
    print(f"F(PK.seed, ADRS, in)={x}")

print(f"wots_pk:{x}")
# bded3fd4b39b5d4a6840c0da0d29aee9
print(f"OK:     bded3fd4b39b5d4a6840c0da0d29aee9")