pure python sha256 implementation and hash-length-extension proof of concept

mysha.py

#!/usr/bin/env python


import struct


k = [
   0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5, 0x3956c25b, 0x59f111f1, 0x923f82a4, 0xab1c5ed5,
   0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3, 0x72be5d74, 0x80deb1fe, 0x9bdc06a7, 0xc19bf174,
   0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc, 0x2de92c6f, 0x4a7484aa, 0x5cb0a9dc, 0x76f988da,
   0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7, 0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967,
   0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, 0x53380d13, 0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85,
   0xa2bfe8a1, 0xa81a664b, 0xc24b8b70, 0xc76c51a3, 0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070,
   0x19a4c116, 0x1e376c08, 0x2748774c, 0x34b0bcb5, 0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3,
   0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208, 0x90befffa, 0xa4506ceb, 0xbef9a3f7, 0xc67178f2
]


def rol(x, n):
    return (((x << 32) | x) >> (32-n)) & 0xFFFFFFFF


def ror(x, n):
    return (((x << 32) | x) >> n) & 0xFFFFFFFF


def add(*values):
    return sum(values) & 0xFFFFFFFF


def bnot(a):
    return (~a) & 0xFFFFFFFF


def sha256_padding(data):
    # Pad out to 512 bits (64 bytes)
    l = len(data) * 8

    # We must leave room for a 1 bit up front (0x80) and then a 64-bit (8-byte) 
    # length field at the end of the padding.
    pad_zeros = (512 - ((l % 512) + 8 + 64)) // 8

    # If the amount of padding is negative we need to add more zeros
    if pad_zeros < 0:
        pad_zeros += 64

    # The length in bits is appended to the padding
    padding = b'\x80' + bytes(pad_zeros) + struct.pack('>Q', l)
    assert (len(data) + len(padding)) % 64 == 0

    return padding


def sha256_hash(state, padded):
    h0, h1, h2, h3, h4, h5, h6, h7 = struct.unpack('>IIIIIIII', state)

    # Chunk into 512 bit chunks
    chunks = [padded[i:i+64] for i in range(0, len(padded), 64)]
    for chunk in chunks:
        # Working buffer (64 32-bit words), copy chunk into working buffer
        w = list(struct.unpack('>16I', chunk))

        # Fill in the rest of the initial working buffer
        for i in range(16, 64):
            s1 = ror(w[i-2], 17) ^ ror(w[i-2], 19) ^ (w[i-2] >> 10)
            s0 = ror(w[i-15], 7) ^ ror(w[i-15], 18) ^ (w[i-15] >> 3)
            w.append(add(s1, w[i-7], s0, w[i-16]))

        assert len(w) == 64

        #print('W:             %08x %08x %08x %08x %08x %08x %08x %08x' % (w[0], w[1], w[2], w[3], w[4], w[5], w[6], w[7]))
        #print('               %08x %08x %08x %08x %08x %08x %08x %08x' % (w[8], w[9], w[10], w[11], w[12], w[13], w[14], w[15]))
        #print('               %08x %08x %08x %08x %08x %08x %08x %08x' % (w[16], w[17], w[18], w[19], w[20], w[21], w[22], w[23]))
        #print('               %08x %08x %08x %08x %08x %08x %08x %08x' % (w[24], w[25], w[26], w[27], w[28], w[29], w[30], w[31]))

        a = h0
        b = h1
        c = h2
        d = h3
        e = h4
        f = h5
        g = h6
        h = h7

        # main loop
        for i in range(64):
            s1 = ror(e, 6) ^ ror(e, 11) ^ ror(e, 25)
            ch = (e & f) ^ (bnot(e) & g)
            temp1 = add(h, s1, ch, k[i], w[i])
            s0 = ror(a, 2) ^ ror(a, 13) ^ ror(a, 22)
            maj = (a & b) ^ (a & c) ^ (b & c)
            temp2 = add(s0, maj)

            h = g
            g = f
            f = e
            e = add(d, temp1)
            d = c
            c = b
            b = a
            a = add(temp1, temp2)

        # Add in the results to the current hash value
        h0 = add(h0, a)
        h1 = add(h1, b)
        h2 = add(h2, c)
        h3 = add(h3, d)
        h4 = add(h4, e)
        h5 = add(h5, f)
        h6 = add(h6, g)
        h7 = add(h7, h)

    digest = struct.pack('>IIIIIIII', h0, h1, h2, h3, h4, h5, h6, h7)

    return digest


def mysha256(data=None, state=None):
    if data is None:
        data = bytes()
    elif isinstance(data, str):
        data = data.encode('latin-1')

    if state is None:
        # Default initial state of a SHA-2 256 digest
        h0 = 0x6a09e667
        h1 = 0xbb67ae85
        h2 = 0x3c6ef372
        h3 = 0xa54ff53a
        h4 = 0x510e527f
        h5 = 0x9b05688c
        h6 = 0x1f83d9ab
        h7 = 0x5be0cd19
        state = struct.pack('>IIIIIIII', h0, h1, h2, h3, h4, h5, h6, h7)

    elif isinstance(state, str):
        state = bytes.fromhex(state)

    #print('Initial state: %s' % state.hex(' ', 4))

    padding = sha256_padding(data)
    padded = data + padding
    #print('Padding:       %s' % (padding.hex(' ', 4)))

    assert len(padded) % 64 == 0

    digest = sha256_hash(state, padded)
    #print('Result:        %s' % digest.hex(' ', 4))

    return digest


def mysha256_extend(data=None, state=None, orig_prefix_len=0, orig_data=None):
    assert data and state and orig_prefix_len > 0 and orig_data

    if isinstance(data, str):
        data = data.encode('latin-1')

    if isinstance(state, str):
        state = bytes.fromhex(state)

    if isinstance(orig_data, str):
        orig_data = orig_data.encode('latin-1')

    # Just generate blank padding, we only need this to generate the correct 
    # padding for the original data.
    orig_data_with_placeholder = bytes(orig_prefix_len) + orig_data
    orig_padding = sha256_padding(orig_data_with_placeholder)
    #print('Old Padding:   %s' % (orig_padding.hex(' ', 4)))

    # Now calculate the new padding
    padding = sha256_padding(orig_data_with_placeholder + orig_padding + data)
    #print('New Padding:   %s' % (padding.hex(' ', 4)))

    # Now do the sha256 calculation on the new data and padding with the new 
    # padding
    padded = data + padding

    assert len(padded) % 64 == 0

    digest = sha256_hash(state, padded)
    #print('Result:        %s' % digest.hex(' ', 4))

    # also return the new data to transmit:
    #   the original data + original padding + new data
    return digest, orig_data + orig_padding + data


if __name__ == '__main__':
    import random
    import hashlib

    print('Test 1:')
    m = 'hello world'
    print('Message:       %s' % m)
    h = hashlib.sha256(m.encode())
    print('Digest:        %s' % h.digest().hex(' ', 4))
    out = mysha256(m)
    print('Result:        %s' % out.hex(' ', 4))
    assert out == h.digest()
    print('PASS!')

    print('\nTest 2:')
    s = random.randbytes(32)
    m = random.randbytes(1024)
    print('Message:       %s + %s' % (s.hex(), m.hex()))
    h = hashlib.sha256(s + m)
    print('Digest:        %s' % h.digest().hex(' ', 4))
    out = mysha256(s + m)
    print('Result:        %s' % out.hex(' ', 4))
    assert out == h.digest()
    print('PASS!')

    print('\nTest 3:')
    s = random.randbytes(32)
    m = random.randbytes(1024)
    print('Message:       %s + %s' % (s.hex(), m.hex()))
    h = hashlib.sha256(s + m)
    print('Digest:        %s' % h.digest().hex(' ', 4))

    a = random.randbytes(32)
    print('\nAppend:        %s' % (a.hex()))
    out, m1 = mysha256_extend(a, state=h.digest(), orig_prefix_len=32, orig_data=m)
    print('Result:        %s' % out.hex(' ', 4))
    h1 = hashlib.sha256(s + m1)
    print('Updated:       %s' % h1.digest().hex(' ', 4))
    assert out == h1.digest()
    print('PASS!')