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

"""SHA-256 functions using hashlib and hmac libraries."""
import hashlib
from hmac import HMAC

# NB installing hmac on Python 3 using pip gives error messages:
#   Collecting hashlib
#     Using cached hashlib-20081119.zip (42 kB)
#     Preparing metadata (setup.py) ... error
#     error: subprocess-exited-with-error
#     × python setup.py egg_info did not run successfully.
#     │ exit code: 1
# but, in fact, it actually works OK.

def SHA256(hexval):
    return hashlib.sha256(bytes.fromhex(hexval)).hexdigest()

def HMAC_SHA256(keyhex, msghex):
    return HMAC(bytes.fromhex(keyhex), bytes.fromhex(msghex), 'sha256').hexdigest()

def MGF1_SHA256(msghex, mlen):
    def _sha256(m: bytes):
        return hashlib.sha256(m).digest()
    z = bytes.fromhex(msghex)
    t = b''
    hlen = len(_sha256(t))  # Get hash length
    for ctr in range(0, (mlen + hlen - 1 // hlen)):
        c = ctr.to_bytes(4)  # default bigendian order
        t += _sha256(z + c)
    return t[:mlen].hex()

if __name__ == '__main__':
    # Basic SHA256 with hex-encoded input
    h = SHA256('616263')  # 'abc' in hex
    print(h)
    # ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad
    assert(h == "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad")

    h = SHA256('')  # hash of empty string
    print(h)
    # e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855
    assert(h == "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855")

    # Raw HMAC with byte input
    h = HMAC(b'Jefe', b'what do ya want for nothing?', 'sha256').hexdigest()
    print(h)  # RFC 4231
    # 5bdcc146bf60754e6a042426089575c75a003f089d2739839dec58b964ec3843
    assert(h == "5bdcc146bf60754e6a042426089575c75a003f089d2739839dec58b964ec3843")

    # HMAC with hex-encoded input
    h = HMAC_SHA256("0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b0b", "4869205468657265")
    print(h)  # RFC 4231
    # b0344c61d8db38535ca8afceaf0bf12b881dc200c9833da726e9376c2e32cff7
    assert(h == "b0344c61d8db38535ca8afceaf0bf12b881dc200c9833da726e9376c2e32cff7")

    # MGF1_SHA256
    h = MGF1_SHA256('3b5c056af3ebba70d4c805380420585562b32410a778f558ff951252407647e3', 34)
    print(h)
    assert(h == "5b7eb772aecf04c74af07d9d9c1c1f8d3a90dcda00d5bab1dc28daecdc86eb87611e")
    # 5b7eb772aecf04c74af07d9d9c1c1f8d3a90dcda00d5bab1dc28daecdc86eb87611e
    h = MGF1_SHA256('', 16)
    print(h)
    # df3f619804a92fdb4057192dc43dd748
    h = MGF1_SHA256('df3f619804a92fdb4057192dc43dd748', 34)