from pwn import process, remote
from base64 import b64encode, b64decode

import hlextend

KEY_LENGTH_BYTES = 16
LOCAL = True

if LOCAL:
    p = process(["python", "main.py"])# Define functions for convenience
else:
    p = remote("2b571e00edfb5dd61b8f7f30-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:"

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: str):
    p.recvuntil(menu_string)
    p.sendline(b"3")
    p.recvuntil(b": ")
    p.sendline(token)

token = register()

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

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

print(f"{claims}")
print(f"{mac}")


injection = "|admin=true".encode("latin1")

sha_forge = hlextend.new('sha1')
forged_msg = sha_forge.extend(injection, claims, KEY_LENGTH_BYTES, mac.decode("latin1"))
print(forged_msg)
forged_mac = sha_forge.hexdigest()

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

print("------------ attacked system ---------")
show_flag(secure_token)
print(p.recvuntil(b":").decode())
p.close()