SHA256 written in pure Python3

sha256.py

# Copyright (c) 2017, 2022 Christian Calderon
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.


def right_rotate(n: int, b: int) -> int:
    """Right rotates a 32 bit word n by b bits."""
    return (n << (32 - b)) | (n >> b) & 0xffffffff


def pad(message: bytes) -> bytes:
    """Pads the message to fit the chunk length of SHA2."""
    L = len(message)*8

    if L%512 > 447:
        K = 1023 - L%512 # 1024 (two chunks) - L%512 - 1
    else:
        K = 511 - L%512 # 512 (one chunk) - L%512 - 1
    
    padding = (1 << K) | L
    return message + padding.to_bytes((K + 1)//8, 'big')


def sha256(message: bytes) -> bytes:
    H = [0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, 0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19]

    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,
    ]

    message = pad(message)

    # collect 512 bit/64 byte chunks into an array
    chunks = []
    while message:
        chunks.append(message[:64])
        message = message[64:]

    for chunk in chunks:
        w = []

        for i in range(16):
            w.append((((((chunk[4*i] << 8) | chunk[4*i+1]) << 8) | chunk[4*i+2]) << 8) | chunk[4*i+3])

        for i in range(16, 64):
            s0 = right_rotate(w[i - 15], 7)^right_rotate(w[i - 15], 18)^(w[i - 15] >> 3)
            s1 = right_rotate(w[i - 2], 17)^right_rotate(w[i - 2], 19)^(w[i - 2] >> 10)
            w.append((w[i - 16] + s0 + w[i - 7] + s1)&0xffffffff)

        a, b, c, d, e, f, g, h = H
        for i in range(64):
            S1 = right_rotate(e, 6)^right_rotate(e, 11)^right_rotate(e, 25)
            ch = ((e&f)^(~e&g))
            temp1 = (h + S1 + ch + k[i] + w[i])&0xffffffff
            S0 = right_rotate(a, 2)^right_rotate(a, 13)^right_rotate(a, 22)
            maj = (a&b)^(a&c)^(b&c)
            temp2 = (S0 + maj)&0xffffffff

            h = g
            g = f
            f = e
            e = (d + temp1)&0xffffffff
            d = c
            c = b
            b = a
            a = (temp1 + temp2)&0xffffffff

        H[0] = (H[0] + a)&0xffffffff
        H[1] = (H[1] + b)&0xffffffff
        H[2] = (H[2] + c)&0xffffffff
        H[3] = (H[3] + d)&0xffffffff
        H[4] = (H[4] + e)&0xffffffff
        H[5] = (H[5] + f)&0xffffffff
        H[6] = (H[6] + g)&0xffffffff
        H[7] = (H[7] + h)&0xffffffff

    return b''.join(h_i.to_bytes(4, 'big') for h_i in H)


def test():
    from binascii import hexlify
    import hashlib
    import os
    test_string = os.urandom(1024)
    print(hashlib.sha256(test_string).hexdigest())
    print(hexlify(sha256(test_string)).decode())

if __name__ == '__main__':
    test()