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

"""SPHINCS+-SHA-256 crypto functions."""

# Either use cryptosyspki or pure hashlib functions
from hashlib_pki import *
#from hashlib_pure import *


# All values passed are hex-encoded

# SPX SHA-256 FUNCTIONS...
def BlockPad(PKseed):
    # Pad PK.seed to 64 bytes with zeros (NB in hex)
    return PKseed + "0" * (128 - len(PKseed))

def F(PKseed, adrs, M):
    return SHA256(BlockPad(PKseed) + adrs + M)[:32]

def H(PKseed, adrs, M1, M2):
    return SHA256(BlockPad(PKseed) + adrs + M1 + M2)[:32]

def PRF(seed, adrs):
    return SHA256(seed + adrs)[:32]

def T_len(PKseed, adrs, M):
    return F(PKseed, adrs, M)

def PRF_msg(SKprf, optrand, M):
    return HMAC_SHA256(SKprf, optrand + M)[:32]

def H_msg(R, PKseed, PKroot, M, m):
    return MGF1_SHA256(SHA256(R + PKseed + PKroot + M), m)


if __name__ == '__main__':
    R = 'b77b5397031e67eb585dba86b10b710b'
    PKseed = 'B505D7CFAD1B497499323C8686325E47'
    PKroot = '4FDFA42840C84B1DDD0EA5CE46482020'
    msg= \
    'D81C4D8D734FCBFBEADE3D3F8A039FAA2A2C9957E835AD55B22E75BF57BB556AC8'
    h_msg = H_msg(R, PKseed, PKroot, msg, 34)
    print(f"H_msg={h_msg}")
    # 5b7eb772aecf04c74af07d9d9c1c1f8d3a90dcda00d5bab1dc28daecdc86eb87611e
    print(f"OK   =5b7eb772aecf04c74af07d9d9c1c1f8d3a90dcda00d5bab1dc28daecdc86eb87611e")