#!/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"

tests = {
    "31d6cfe0d16ae931b73c59d7e0c089c0": b"",
    "bde52cb31de33e46245e05fbdbd6fb24": b"a",
    "a448017aaf21d8525fc10ae87aa6729d": b"abc",
    "d9130a8164549fe818874806e1c7014b": b"message digest",
    "d79e1c308aa5bbcdeea8ed63df412da9": b"abcdefghijklmnopqrstuvwxyz",
    "043f8582f241db351ce627e153e7f0e4": b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789",
    "e33b4ddc9c38f2199c3e7b164fcc0536": b"12345678901234567890123456789012345678901234567890123456789012345678901234567890",
    "1bee69a46ba811185c194762abaeae90": b"The quick brown fox jumps over the lazy dog",
    "b86e130ce7028da59e672d56ad0113df": b"The quick brown fox jumps over the lazy cog",
}


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 md4:
    digest_size = 16  # MD4's digest byte size
    block_size = 64  # MD4's block size in bytes

    word_size = 4  # MD4's word size in bytes
    word_bits = word_size * 8  # MD4'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)
        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, "little"))

        # 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))

    def _f(self, x, y, z):
        return (x & y) | (~x & z)

    def _ff(self, a, b, c, d, x, s):
        return self._left_rotate((a + self._f(b, c, d) + x) & self.word_mask, s)

    def _g(self, x, y, z):
        return (x & y) | (x & z) | (y & z)

    def _gg(self, a, b, c, d, x, s):
        return self._left_rotate(
            (a + self._g(b, c, d) + x + 0x5A827999) & self.word_mask, s
        )

    def _h(self, x, y, z):
        return x ^ y ^ z

    def _hh(self, a, b, c, d, x, s):
        return self._left_rotate(
            (a + self._h(b, c, d) + x + 0x6ED9EBA1) & self.word_mask, s
        )

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

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

        S11 = 3
        S12 = 7
        S13 = 11
        S14 = 19
        S21 = 3
        S22 = 5
        S23 = 9
        S24 = 13
        S31 = 3
        S32 = 9
        S33 = 11
        S34 = 15

        # Step 1
        a = self._ff(a, b, c, d, w[0], S11)
        d = self._ff(d, a, b, c, w[1], S12)
        c = self._ff(c, d, a, b, w[2], S13)
        b = self._ff(b, c, d, a, w[3], S14)
        a = self._ff(a, b, c, d, w[4], S11)
        d = self._ff(d, a, b, c, w[5], S12)
        c = self._ff(c, d, a, b, w[6], S13)
        b = self._ff(b, c, d, a, w[7], S14)
        a = self._ff(a, b, c, d, w[8], S11)
        d = self._ff(d, a, b, c, w[9], S12)
        c = self._ff(c, d, a, b, w[10], S13)
        b = self._ff(b, c, d, a, w[11], S14)
        a = self._ff(a, b, c, d, w[12], S11)
        d = self._ff(d, a, b, c, w[13], S12)
        c = self._ff(c, d, a, b, w[14], S13)
        b = self._ff(b, c, d, a, w[15], S14)

        # Step 2
        a = self._gg(a, b, c, d, w[0], S21)
        d = self._gg(d, a, b, c, w[4], S22)
        c = self._gg(c, d, a, b, w[8], S23)
        b = self._gg(b, c, d, a, w[12], S24)
        a = self._gg(a, b, c, d, w[1], S21)
        d = self._gg(d, a, b, c, w[5], S22)
        c = self._gg(c, d, a, b, w[9], S23)
        b = self._gg(b, c, d, a, w[13], S24)
        a = self._gg(a, b, c, d, w[2], S21)
        d = self._gg(d, a, b, c, w[6], S22)
        c = self._gg(c, d, a, b, w[10], S23)
        b = self._gg(b, c, d, a, w[14], S24)
        a = self._gg(a, b, c, d, w[3], S21)
        d = self._gg(d, a, b, c, w[7], S22)
        c = self._gg(c, d, a, b, w[11], S23)
        b = self._gg(b, c, d, a, w[15], S24)

        # Step 3
        a = self._hh(a, b, c, d, w[0], S31)
        d = self._hh(d, a, b, c, w[8], S32)
        c = self._hh(c, d, a, b, w[4], S33)
        b = self._hh(b, c, d, a, w[12], S34)
        a = self._hh(a, b, c, d, w[2], S31)
        d = self._hh(d, a, b, c, w[10], S32)
        c = self._hh(c, d, a, b, w[6], S33)
        b = self._hh(b, c, d, a, w[14], S34)
        a = self._hh(a, b, c, d, w[1], S31)
        d = self._hh(d, a, b, c, w[9], S32)
        c = self._hh(c, d, a, b, w[5], S33)
        b = self._hh(b, c, d, a, w[13], S34)
        a = self._hh(a, b, c, d, w[3], S31)
        d = self._hh(d, a, b, c, w[11], S32)
        c = self._hh(c, d, a, b, w[7], S33)
        b = self._hh(b, c, d, a, w[15], S34)

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

    def digest(self):
        return struct.pack("<4I", *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("<4I", b)


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


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


def create_md4_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, "little"))
    return chunk


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

        h = md4()
        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


def test_md4():
    for hash, data in tests.items():
        h = md4()
        h.update(data)
        my_hash = h.hexdigest()
        if my_hash != hash:
            print(f"Failed on: {my_hash} != {hash} for {data}")
            return False

        h = md4()
        for l in data:
            h.update(bytes([l]))
        my_hash = h.hexdigest()
        if my_hash != hash:
            print(f"Failed on multi update: {my_hash} != {hash} for {data}")
            return False
    return True


if test_md4():
    print("MD4 Tests Passed!")
else:
    print("MD4 Tests Failed!")

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 :(")
