A Simple SHA256 Implementation in python

SHA256.py

from numpy import uint32
from numpy import uint64
from numpy import geterr
from numpy import seterr


class sha256:
    STARTING_VALUES = [
        uint32(0x6a09e667),
        uint32(0xbb67ae85),
        uint32(0x3c6ef372),
        uint32(0xa54ff53a),
        uint32(0x510e527f),
        uint32(0x9b05688c),
        uint32(0x1f83d9ab),
        uint32(0x5be0cd19)
    ]

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

    def __init__(self, message):
        self.h = self.STARTING_VALUES.copy()
        self.message = bytearray(message)

    def reset(self):
        self.h = self.STARTING_VALUES.copy()
        self.message = bytearray()

    def update(self, message_bytes):
        self.message += bytes(message_bytes)

    def _digest_chunk(self, chunk):

        def rrotate(input, r):
            return (uint32(input) >> uint32(r)) | \
                   (uint32(input) << uint32(32 - r))

        if len(chunk) == 64:
            w = []
            for i in range(0, 64, 4):
                w.append(uint32(int.from_bytes(chunk[i:i + 4], 'big')))

            for i in range(16, 64):
                w.append(w[i-16] + w[i-7])

                t0 = rrotate(w[i-15], 7)
                t0 ^= rrotate(w[i-15], 18)
                t0 ^= w[i-15] >> uint32(3)
                w[i] += t0

                t1 = rrotate(w[i-2], 17)
                t1 ^= rrotate(w[i-2], 19)
                t1 ^= w[i-2] >> uint32(10)
                w[i] += t1

            d = self.h.copy()

            for i in range(0, 64):
                a1 = rrotate(d[4], 6)
                a1 ^= rrotate(d[4], 11)
                a1 ^= rrotate(d[4], 25)
                b1 = (d[4] & d[5]) ^ (~d[4] & d[6])
                temp1 = d[7] + a1 + b1 + self.K_VALUES[i] + w[i]

                a2 = rrotate(d[0], 2)
                a2 ^= rrotate(d[0], 13)
                a2 ^= rrotate(d[0], 22)
                b2 = (d[0] & d[1]) ^ (d[0] & d[2]) ^ (d[1] & d[2])
                temp2 = a2 + b2

                d[7] = d[6]
                d[6] = d[5]
                d[5] = d[4]
                d[4] = d[3] + temp1
                d[3] = d[2]
                d[2] = d[1]
                d[1] = d[0]
                d[0] = temp1 + temp2

            self.h[0] += d[0]
            self.h[1] += d[1]
            self.h[2] += d[2]
            self.h[3] += d[3]
            self.h[4] += d[4]
            self.h[5] += d[5]
            self.h[6] += d[6]
            self.h[7] += d[7]

    def digest(self):

        def to_bytes_64(uint64_input):
            output = bytearray()
            for i in range(0, 8):
                output.extend(bytearray([uint64_input & uint64(0xff)]))
                uint64_input >>= uint64(8)
            return bytes(reversed(output))

        def to_bytes_32(uint32_arr):
            output = bytearray()
            for i in range(0, 4):
                output.extend(bytearray([uint32_arr & uint32(0xff)]))
                uint32_arr >>= uint32(8)
            return bytes(reversed(output))

        # Turn off overflow errors for numpy
        temp_set = geterr()['over']
        seterr(over='ignore')

        # Get a copy of the message to digest
        message = self.message.copy()

        # Get message length in bits
        message_len = uint64(len(self.message) * 8)

        # Add set bit to end of message
        self.message += bytearray([0x80])

        # Add padding 0's to end of message
        self.message += bytearray([0] * (64 - (len(self.message) + 8) % 64))

        # Add message length to message
        self.message += to_bytes_64(message_len)

        # Reset H values for digest
        self.h = self.STARTING_VALUES.copy()

        # Digest the message in chunks
        for i in range(0, len(self.message), 64):
            self._digest_chunk(self.message[i:i + 64])

        # Revert numpy error settings
        seterr(over=temp_set)

        # Return digest
        digest = bytearray()
        for word in self.h:
            digest += to_bytes_32(word)
        return bytes(digest)

    def hex_digest(self):
        digest = self.digest()
        return digest.hex()