from Crypto.Cipher import AES
from Crypto.Util.strxor import strxor
import struct

class SpongeBob():
    def __init__(self):
        self.key = b'NOT_RLY_A_SECRET'
        self.state = b'\x00'*16
        self.bitrate_bytes = 10
        self.aes = AES.new(self.key, AES.MODE_ECB)

    def __pad(self, data):
        # data + padding + 4B_length
        data_size = len(data)
        padding_size = self.bitrate_bytes - ((data_size + 4) % self.bitrate_bytes)
        data += struct.pack('B', padding_size) * padding_size
        data += struct.pack('>I', data_size)
        return data

    def __absorb(self, data):
        # Feed sponge with data in chunks
        chunks = [data[i:i+self.bitrate_bytes] for i in range(0, len(data), self.bitrate_bytes)]
        for chunk in chunks:
            chunk = bytes(chunk) + b'\x00' * 6
            block = strxor(chunk, self.state)
            self.state = self.aes.encrypt(block)

    def __squeeze(self):
        result = self.state[:self.bitrate_bytes]
        self.state = self.aes.encrypt(self.state)
        return result

    def hash(self, data):
        data = self.__pad(data)
        self.__absorb(data)
        return self.__squeeze() + self.__squeeze() + self.__squeeze()


def main():
    print('Start')
    message = b'The hash-slinging slasher'
    bob = SpongeBob()
    h = bob.hash(message)
    print(f'Message: {message}')
    print(f'Hash: {h.hex()}')



if __name__ == '__main__':
    main()