# @file wots_sig_0.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 first WOTS signature."""

from spx_adrs import Adrs
from spx_sha256 import PRF, F

PKseed = 'B505D7CFAD1B497499323C8686325E47'
SKseed = '7c9935a0b07694aa0c6d10e4db6b1add'

# Working in the bottom HT subtree...
# Set up ADRS object
adrs = Adrs(Adrs.WOTS_HASH, layer=0)
adrs.setTreeAddress(0x28daecdc86eb8761)
adrs.setKeyPairAddress(6)
print(f"ADRS base={adrs.toHex()}")

idx = 0
print(f"Generate WOTS+ private key for i = {idx}")
# sk = PRF(SK.seed, ADRS)
adrs.setChainAddress(idx)
adrs_c = adrs.toHex()
print(f"ADRS={adrs_c}")
sk = PRF(SKseed, adrs_c)
print(f"sk={sk}")

# Compute F^m(sk)
m = 0xd  # m_0
print(f"m[{idx}]={m}")
x = sk
adrs_ht = Adrs.fromHex(adrs.toHex())
for i in range(m):
    adrs_ht.setHashAddress(i)
    adrs_c = adrs_ht.toHex()
    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"ht_sig:{x}")
print("ht_sig[0][0]:cb56a9488bb970ce78a2623db324b873")
assert(x == "cb56a9488bb970ce78a2623db324b873")

print("")
 
idx = 34
print(f"Generate WOTS+ private key for i = {idx}")
# sk = PRF(SK.seed, ADRS)
adrs.setChainAddress(idx)
adrs_c = adrs.toHex()
print(f"ADRS={adrs_c}")
sk = PRF(SKseed, adrs_c)
print(f"sk={sk}")

# Compute F^m(sk)
m = 0xf  # m_34
print(f"m[{idx}]={m}")
x = sk
adrs_ht = Adrs.fromHex(adrs.toHex())
for i in range(m):
    adrs_ht.setHashAddress(i)
    adrs_c = adrs_ht.toHex()
    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"ht_sig:{x}")
print("ht_sig[0][34]:a0154e91a720b96b0d738d05405e6f48") 
assert(x == "a0154e91a720b96b0d738d05405e6f48")