# @file wots_pk.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 the WOTS public key for the bottom HT subtree."""

from spx_adrs import Adrs
from spx_sha256 import H, F, PRF, T_len
from spx_util import chain

PKseed = 'B505D7CFAD1B497499323C8686325E47'
SKseed = '7c9935a0b07694aa0c6d10e4db6b1add'
w = 16

# Working in the bottom HT subtree...
wots_pk_hex = ""
for idx in range(35):
    print(f"Generate WOTS+ private key for subtree[{idx}]")
    # sk = PRF(SK.seed, ADRS)
    adrs = Adrs(Adrs.WOTS_HASH, layer=0)
    adrs.setTreeAddress(0x28daecdc86eb8761)
    adrs.setKeyPairAddress(6)
    adrs.setChainAddress(idx)
    print(f"ADRS={adrs.toHex()}")
    sk = PRF(SKseed, adrs.toHex())
    print(f"sk[{idx}]={sk}")
    if idx == 0:
        print(f"OK   =c04623124dfcdcb1de0ad8cfc68ebf73")
    elif idx == 1:
        print(f"OK   =5ae611d347f988308958c414f865c942")

    print(f"Compute pk...")
    # Compute F^{w-1}(sk)
    # Note that ADRS is the same for sk and pk chaining
    adrs = Adrs(Adrs.WOTS_HASH, layer=0)
    adrs.setTreeAddress(0x28daecdc86eb8761)
    adrs.setKeyPairAddress(6)
    adrs.setChainAddress(idx)
    print(f"ADRS={adrs.toHex()}")
    pk = chain(sk, 0, w - 1, PKseed, adrs.toHex(), True if idx == 0 else False)
    print(f"pk[{idx}]={pk}")
    if idx == 0:
        print(f"OK   =aff586153b484118f05accd6e7bde2a3")
    elif idx == 1:
        print(f"OK   =0a7e84c64a9d47f3a9ec13092c2b68dc")
    elif idx == 34:
        print(f"OK    =a0154e91a720b96b0d738d05405e6f48")
    wots_pk_hex += pk

print(f"wots_pk={wots_pk_hex}")


# We have all 35 wots_pk values for the first subtree
# Compress wots_pk values to n-byte leaf value

# HARDCODED - not used
wots_pk="aff586153b484118f05accd6e7bde2a30a7e84c64a9d47f3a9ec13092c2b68dcff648dce82a1e3602c736bae67fc0ec8b456ee1c1f4a886fd34b07a30c99080c7f8a54ce8028be337b4acb99df5e7c119be45116572eeb1b2f4e267785d76c02581eae8a317beab136fec07e6fe56c39ca8761a4844d2735177f85702781466fa9276f2c973cb290a566abdf6942d4c1534ddbe65e6c006f5b6de33c7d42d1386addf4ec3bd6b8ec2195f9b56639fbeaa839015a44a2d6fda22e184f0c1b47a2c4e960d263ff1a13cc05b8816fd9926f101fbb0a60498cfe239afd6b5e873645fcb2d895191a5a745136b0b3eca8155b96fb21427d635f6555ee14878743ee7f2b014c0cf0d63a5c4ff6ddeca4e65b165bb93f17219b9f7e8ed1207ad186beb0f724c8e67b28bff9164391fd2370d1f3c3f3d8cfbdbe0f4b13034efb6e317738b1520bee600c6595aed066ff513be6d6b7bb780d7e3f7a09616a4803c53e4e5e7ebb6b03d6a96f02f6fad42a6a56291f30b1deb762aaf573313eba8bb31f6efb3ffdca706c90c07533c124ac91754b3e18234c078e8c2e32d5ea1e97b9426aceb523382686ce6313e56bb20badff38c76b0b66100d08832a5d670e83c32334127f2adf6faa0aaa3ff8b4d0b1efecc4b13564e895d7b836670e5b30a860a6e56b30f8ac6f402346901f8cbbdb74a9c83ffb4e558e134d9ff15ecdcaa132d59edd8d89023eac07a65ceede697aaba750271377cfa250e29aec2dac2ada65ae3d98a0154e91a720b96b0d738d05405e6f48"
wots_addr="0028daecdc86eb876101000000060000000000000000"
print(f"ADRS={wots_addr}")

# Computed values...
adrs = Adrs(Adrs.WOTS_PK, layer=0)
adrs.setTreeAddress(0x28daecdc86eb8761)
adrs.setKeyPairAddress(6)
print(f"ADRS={adrs.toHex()}")
leaf = T_len(PKseed, adrs.toHex(), wots_pk_hex)
print(f"leaf={leaf}")
# leaf=thash(wots_pk)=56c1cd468c05d6b5a9ad57e87c4edf12
print(f"OK  =56c1cd468c05d6b5a9ad57e87c4edf12")
assert leaf == "56c1cd468c05d6b5a9ad57e87c4edf12"