from pwn import *
from base64 import b64encode, b64decode
import struct
import binascii

KEY_LENGTH_BYTES = 16
LOCAL = True

if not LOCAL:
    p = remote("190b6910a0ded80a0296d180-1024-intro-crypto-2.challenge.cscg.live", 1337, ssl=True)# Define functions for convenience

menu_string = b"\n        1. Register\n        2. Show animal videos\n        3. Show flag\n        4. Exit\n        \nEnter your choice:"

class SHA1:
    def __init__(self, h0, h1, h2, h3, h4, message_byte_length):
        self.h0 = h0
        self.h1 = h1
        self.h2 = h2
        self.h3 = h3
        self.h4 = h4
        self.message_byte_length = message_byte_length

    def _left_rotate(self, n, b):
        return ((n << b) | (n >> (32 - b))) & 0xffffffff

    def update(self, message):
        message_byte_length = len(message)
        message_bit_length = (self.message_byte_length + message_byte_length) * 8
        message += b'\x80'
        message += b'\x00' * ((56 - (len(message) % 64)) % 64)
        message += struct.pack('>Q', message_bit_length)

        for chunk_offset in range(0, len(message), 64):
            w = list(struct.unpack('>16I', message[chunk_offset:chunk_offset + 64]))
            w.extend(0 for _ in range(64))
            for i in range(16, 80):
                w[i] = self._left_rotate(w[i - 3] ^ w[i - 8] ^ w[i - 14] ^ w[i - 16], 1)

            a, b, c, d, e = self.h0, self.h1, self.h2, self.h3, self.h4
            for i in range(80):
                if 0 <= i < 20:
                    f = (b & c) | (~b & d)
                    k = 0x5A827999
                elif 20 <= i < 40:
                    f = b ^ c ^ d
                    k = 0x6ED9EBA1
                elif 40 <= i < 60:
                    f = (b & c) | (b & d) | (c & d)
                    k = 0x8F1BBCDC
                else:
                    f = b ^ c ^ d
                    k = 0xCA62C1D6

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

            self.h0 = (self.h0 + a) & 0xffffffff
            self.h1 = (self.h1 + b) & 0xffffffff
            self.h2 = (self.h2 + c) & 0xffffffff
            self.h3 = (self.h3 + d) & 0xffffffff
            self.h4 = (self.h4 + e) & 0xffffffff

    def digest(self):
        return struct.pack('>5I', self.h0, self.h1, self.h2, self.h3, self.h4)

def recv_until(delim):
    return p.recvuntil(delim)

def interact():
    p.interactive()

def send_payload(payload):
    p.sendline(payload)

def register() -> bytes:
    p.recvuntil(menu_string)
    p.sendline(b"1")
    p.recvuntil(b"? ")
    p.sendline(b"Simon")
    p.recvuntil(b"? ")
    p.sendline(b"Cat")
    return p.recvuntil(b'\n').replace(b"\n", b"").split(b" ")[-1]

def show_flag(token: bytes):
    p.recvuntil(menu_string)
    p.sendline(b"3")
    p.recvuntil(b": ")
    p.sendline(token)
    interact()

def sha1_padding(message_length):
    padding = b'\x80'
    padding += b'\x00' * ((56 - (message_length + 1) % 64) % 64)
    padding += struct.pack('>Q', message_length * 8)
    return padding

def get_mac(data: bytes) -> str:
    return sha1(KEY.encode("latin1") + data).hexdigest()


if LOCAL:
    print("Enter token:")
    token = input()
else:
    token = register()

string_token = b64decode(token)
claims, mac = string_token.split(b"|mac=")

print(f"{token} : {string_token}")

padding = sha1_padding(KEY_LENGTH_BYTES + len(claims))
injection = "|admin=true".encode("latin1")

forged_msg = claims + padding + injection 

h = struct.unpack('>5I', bytes.fromhex(mac.decode("latin1")))

sha1 = SHA1(*h, message_byte_length=len(forged_msg) - len(injection))
sha1.update(injection)

forged_mac = sha1.digest().hex()

key = input("Input Key")

orig_sha1 = SHA1(*[0x67452301, 0xEFCDAB89, 0x98BADCFE, 0x10325476, 0xC3D2E1F0], message_byte_length=(len(forged_msg) + KEY_LENGTH_BYTES))
orig_sha1.update(key.encode("latin1") + forged_msg)

real_mac = orig_sha1.digest().hex()


print(f"forged mac: {forged_mac} | real mac: {real_mac}")

secure_token = forged_msg + f"|mac={forged_mac}".encode("latin1")
secure_token = b64encode(secure_token).decode("latin1")

if LOCAL:
    print(secure_token)
else:
    show_flag(secure_token)