# @file H_msg.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 H_msg from the message, public key and randomizer."""

from spx_sha256 import SHA256, MGF1_SHA256

# digest = H_Msg(R, PK.seed, PK.root, M);
# = MGF1-SHA-256(SHA-256(R||PK.seed||PK.root||M), m)
# m = SPX_DGST_BYTES = 34

R = 'b77b5397031e67eb585dba86b10b710b'
PKseed = 'B505D7CFAD1B497499323C8686325E47'
PKroot = '4FDFA42840C84B1DDD0EA5CE46482020'
msg= \
'D81C4D8D734FCBFBEADE3D3F8A039FAA2A2C9957E835AD55B22E75BF57BB556AC8'

print("Input =", R + PKseed + PKroot + msg)
h = SHA256(R + PKseed + PKroot + msg)
print("h =", h)
# 3b5c056af3ebba70d4c805380420585562b32410a778f558ff951252407647e3
h_msg = MGF1_SHA256(h, 34)
print("H_msg:", h_msg)
# 5b7eb772aecf04c74af07d9d9c1c1f8d3a90dcda00d5bab1dc28daecdc86eb87611e

# Split into mhash (25 bytes), tree address (8 bytes) and idx_leaf (1 byte)
print("Split up H_msg...")
mhash = h_msg[:50]
tree_hex = h_msg[50:66]
leaf_hex = h_msg[66:68]
print(f"mhash='{mhash}'")
print(f"tree='{tree_hex}', leaf='{leaf_hex}'")
# Decode tree address and leaf index into integers
tree_addr = int(tree_hex, 16) & 0x7fffffffffffffff  # 63 bits
print("tree_addr = 0x" + format(tree_addr, f'08x'))
idx_leaf = int(leaf_hex, 16) & 0x7  # 3 bits
print(f"idx_leaf = {idx_leaf}")