from pwn import context, args, process, remote, success, xor
from base64 import b64decode, b64encode
from hashlib import sha1
import struct

context.update(arch="amd64", os="linux", log_level="info")

URL = "jpielr6wv7xor5jf32vxasubni-1024-intro-crypto-2.challenge.cscg.live"
PORT = 443

example_token="name=cato|animal=cat|admin=false|mac=b5d4ddd83235359176d107b7c98bc7f2979cfec3"


class Sha1State:
    def __init__(self, hex_digest):
        # Extract the 5 internal 32-bit registers from the hex digest
        self.h = struct.unpack(">5I", bytes.fromhex(hex_digest))

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

    def extend(self, data, total_len_bytes):
        """
        Continues hashing from the current state. 
        total_len_bytes must be the length of (secret + message + padding + extra_data)
        """
        h0, h1, h2, h3, h4 = self.h
        
        # SHA-1 processes 64-byte blocks
        for i in range(0, len(data), 64):
            block = data[i:i+64]
            if len(block) < 64:
                # This is the final block of the 'extra_data', it needs its own padding
                block += self.get_padding(total_len_bytes)
            
            w = list(struct.unpack(">16I", block[:64]))
            for j in range(16, 80):
                w.append(self._left_rotate(w[j-3] ^ w[j-8] ^ w[j-14] ^ w[j-16], 1))

            a, b, c, d, e = h0, h1, h2, h3, h4

            for j in range(80):
                if 0 <= j <= 19:
                    f = (b & c) | ((~b) & d)
                    k = 0x5A827999
                elif 20 <= j <= 39:
                    f = b ^ c ^ d
                    k = 0x6ED9EBA1
                elif 40 <= j <= 59:
                    f = (b & c) | (b & d) | (c & d)
                    k = 0x8F1BBCDC
                elif 60 <= j <= 79:
                    f = b ^ c ^ d
                    k = 0xCA62C1D6

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

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

        return '%08x%08x%08x%08x%08x' % (h0, h1, h2, h3, h4)

    def get_padding(self, msg_len):
        padding = b"\x80"
        padding += b"\x00" * ((56 - (msg_len + 1)) % 64)
        padding += struct.pack(">Q", msg_len * 8)
        return padding

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

def forge_mac(orig_mac: str, key_len: int, orig_data: bytes, added_data: bytes) -> str:
    state = Sha1State(orig_mac)
    current_len = key_len + len(orig_data)
    glue_padding = get_sha1_padding(current_len)
    
    new_total_len = current_len + len(glue_padding) + len(added_data)
    
    return state.extend(added_data, new_total_len)

def start():
    if args.REMOTE:
        return remote(URL, PORT, ssl=True)
    else:
        return process(["python3", "./main.py"])


def get_token(io):
    io.sendline(b"1")
    io.recvuntil(b"? ")
    io.sendline(b"cato")
    io.recvuntil(b"? ")
    io.sendline(b"cat")
    return io.recvline().split(b" ")[-1].strip()

def get_flag(io, token):
    io.sendline(b"3")
    io.recvuntil(b": ")
    io.sendline(token)
    return io.recvline()


io = start()
io.recvuntil(b": ")
token = b64decode(get_token(io))
parts = token.split(b"|mac=")
data_to_pad = parts[0]
original_mac = parts[1].decode()

extension = b"|admin=true"
key_len = 32
glue = get_sha1_padding(key_len + len(data_to_pad))
new_mac = forge_mac(original_mac, key_len, data_to_pad, extension)
malicious_payload = data_to_pad + glue + extension + b"|mac=" + new_mac.encode()
final_token = b64encode(malicious_payload)
io.recvuntil(b": ")
print(get_flag(io, final_token).decode())