#!/usr/bin/env python3
import sys
import struct
from random import randint

input = b"comment1=cooking%20MCs;userdata=foo;comment2=%20like%20a%20pound%20of%20bacon"


def random_bytes(size):
    i = randint(0, 2 ** (8 * size))
    return i.to_bytes(size, sys.byteorder)


unknown_key = random_bytes(randint(16, 64))


class sha1:
    digest_size = 20  # SHA1's digest byte size
    block_size = 64  # SHA1's block size in bytes

    word_size = 4  # SHA1's word size in bytes
    word_bits = word_size * 8  # SHA1's word size in bits
    word_mask = 2 ** 32 - 1  # Bitwise AND mask to limit int's to word_size

    def __init__(self):
        self._state = (0x67452301, 0xEFCDAB89, 0x98BADCFE, 0x10325476, 0xC3D2E1F0)
        self._message_size = 0  # Message size in bytes
        self._pending = bytearray()

    def _get_final_state(self):
        chunk = bytearray(self._pending)

        # The final real message size
        message_size = self._message_size + len(chunk)

        # Pre-processing to pad the message to a multiple of block_size
        chunk.append(0x80)
        # Magic Number 9: 1 byte for above pad & 8 bytes for final message size suffix
        remaining_length = (
            self.block_size - ((message_size + 9) % self.block_size)
        ) % self.block_size  # On perfect message size add no padding
        chunk.extend(b"\x00" * remaining_length)

        # Append the real message_size in bits as a 64bit int
        chunk.extend((message_size * 8).to_bytes(8, "big"))

        # Padded chunk maybe 2 blocks worth
        if len(chunk) > self.block_size:
            state = self._process_chunk(chunk[: self.block_size], self._state)
            chunk = chunk[self.block_size :]
        else:
            state = self._state
        return self._process_chunk(chunk, state)

    def _left_rotate(self, number, amount):
        return (
            (number << amount) | (number >> (self.word_bits - amount))
        ) & self.word_mask

    def _get_words(self, chunk):
        w = [
            int.from_bytes(chunk[i : i + self.word_size], "big")
            for i in range(0, len(chunk), self.word_size)
        ]

        for i in range(16, 80):
            w.append(
                self._left_rotate((w[i - 3] ^ w[i - 8] ^ w[i - 14] ^ w[i - 16]), 1)
            )
        return w

    def _process_chunk(self, chunk, state):
        a, b, c, d, e = state
        w = self._get_words(chunk)

        for i in range(0, 80):
            if i <= 19:
                f = d ^ (b & (c ^ d))
                k = 0x5A827999
            elif 20 <= i and i <= 39:
                f = b ^ c ^ d
                k = 0x6ED9EBA1
            elif 40 <= i and i <= 59:
                f = (b & c) ^ (d & (b ^ c))
                k = 0x8F1BBCDC
            else:
                f = b ^ c ^ d
                k = 0xCA62C1D6

            temp = ((self._left_rotate(a, 5)) + f + e + k + w[i]) & self.word_mask
            e = d
            d = c
            c = self._left_rotate(b, 30)
            b = a
            a = temp

        return (
            (state[0] + a) & self.word_mask,
            (state[1] + b) & self.word_mask,
            (state[2] + c) & self.word_mask,
            (state[3] + d) & self.word_mask,
            (state[4] + e) & self.word_mask,
        )

    def digest(self):
        return struct.pack(">5I", *self._get_final_state())

    def hexdigest(self):
        return self.digest().hex()

    def update(self, data):
        # Create the first block to update with based on the queued data in _pending
        next = bytearray(self._pending)
        i = self.block_size - len(next)
        next.extend(data[:i])

        # Loop though all full blocks and process them
        while len(next) == self.block_size:
            self._state = self._process_chunk(next, self._state)
            self._message_size += self.block_size
            next = data[i : i + self.block_size]
            i += self.block_size

        # What ever remains afterwards is put into _pending for later
        self._pending = bytearray(next)

    def extract_state_hex(self, s, length):
        self.extract_state(bytes.fromhex(s), length)

    def extract_state(self, b, length):
        self._message_size = length
        self._state = struct.unpack(">5I", b)


def get_mac(data):
    h = sha1()
    h.update(unknown_key)
    h.update(data)
    return h.hexdigest()


def validate_mac(data, mac):
    return get_mac(data) == mac


def create_sha1_padding(message_size):
    # Pre-processing to pad the message to a multiple of block_size
    chunk = bytearray([0x80])

    # Magic Number 9: 1 byte for above pad & 8 bytes for final message size suffix
    remaining_length = 64 - ((message_size + 9) % 64)
    chunk.extend(b"\x00" * remaining_length)

    # Append the real message_size in bits as a 64bit int
    chunk.extend((message_size * 8).to_bytes(8, "big"))
    return chunk


def find_valid_extension_mac(data, orig_mac, extension):
    data_length = len(input)
    for i in range(0, 128):
        padding = create_sha1_padding(data_length + i)
        padded_size = data_length + len(padding) + i

        h = sha1()
        h.extract_state_hex(orig_mac, padded_size)
        h.update(extension)
        mac = h.hexdigest()

        target = input + padding + extension
        if validate_mac(target, mac):
            return target, mac
    return None, None


new_data, new_mac = find_valid_extension_mac(input, get_mac(input), b";admin=true")
if new_data is not None and new_mac is not None:
    print(f"Got Admin!: {new_data} with mac {new_mac}")
else:
    print("Not an Admin this day :(")
