import socket
import base64
import hashlib
from Crypto.Cipher import AES
import binascii
# Fill in the right target here
HOST = 'netsec.net.in.tum.de'  # TODO
PORT = 20105  # TODO
KEY = b'1337133713371337'

def pkcs7(message: bytes, block_size: int = 16) -> bytes:
    gap_size = block_size - (len(message) % block_size)
    return message + bytes([gap_size] * gap_size)

def calc_hmac(key: bytes, message: bytes ) -> bytes:
    block_size = 64 
    if len(key) > block_size:
        key = hashlib.sha256(key).digest()
    if len(key) < block_size:
        key += b'\x00' * (block_size - len(key))

    o_key_pad = bytes([k ^ 0x5C for k in key])
    i_key_pad = bytes([k ^ 0x36 for k in key])

    inner_hash = hashlib.sha256(i_key_pad + message).digest()
    return hashlib.sha256(o_key_pad + inner_hash).digest()

def calc_cbc_mac( key: bytes, iv: bytes, message: bytes) -> bytes:
    message = pkcs7(message)
    cipher = AES.new(key, AES.MODE_CBC, iv)
    encrypted_message = cipher.encrypt(message)
    return encrypted_message[-16:]

def cmac_padding(message: bytes, block_size: int = 16) -> bytes:
    if len(message) % block_size == 0:
        return message  # Full block, no padding needed
    else:
        padded_message = message + b'\x80'  # Add the 0x80 byte
        return padded_message.ljust((len(message) // block_size + 1) * block_size, b'\x00')

def calc_cmac(key: bytes, message: bytes) -> bytes:
    from Crypto.Cipher import AES

    def shift_left(block: bytes) -> bytes:
        result = int.from_bytes(block, byteorder="big") << 1
        return (result & ((1 << 128) - 1)).to_bytes(16, byteorder="big")

    const_rb = 0x87
    cipher = AES.new(key, AES.MODE_ECB)
    zero_block = b'\x00' * 16
    l_block = cipher.encrypt(zero_block)

    k1 = shift_left(l_block)
    if l_block[0] & 0x80:
        k1 = (int.from_bytes(k1, byteorder="big") ^ const_rb).to_bytes(16, byteorder="big")

    k2 = shift_left(k1)
    if k1[0] & 0x80:
        k2 = (int.from_bytes(k2, byteorder="big") ^ const_rb).to_bytes(16, byteorder="big")

    padded_message = cmac_padding(message)
    if len(message) % 16 == 0:
        last_block = bytes([b ^ k1[i] for i, b in enumerate(padded_message[-16:])])
    else:
        last_block = bytes([b ^ k2[i] for i, b in enumerate(padded_message[-16:])])

    previous_block = zero_block
    for i in range(0, len(padded_message) - 16, 16):
        block = padded_message[i:i + 16]
        previous_block = cipher.encrypt(bytes([b ^ p for b, p in zip(block, previous_block)]))

    cmac_result = cipher.encrypt(bytes([b ^ p for b, p in zip(last_block, previous_block)]))
    return cmac_result

def decode_message(msg):
    return binascii.unhexlify(msg)


def get_flag():
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)

    s.connect((HOST, PORT))
    sf = s.makefile('rw')  # we use a file abstraction for the sockets

    message1 = sf.readline().rstrip('\n')
    print(message1)
    challenge = decode_message(message1)
    hmac = base64.b64encode(calc_hmac(KEY, challenge)).decode()
    cmac = base64.b64encode(calc_cmac(KEY, challenge)).decode()
    cbc_mac = base64.b64encode(calc_cbc_mac(KEY,  b'\x00' * 16, challenge)).decode()
    answer = f"{hmac};{cbc_mac};{cmac}"
    print(f"Calcualted the answer {answer}")
    sf.write(f"{answer}\n")
    sf.flush()
    print(sf.readline().rstrip('\n'))
    sf.close()
    s.close()


if __name__ == '__main__':
    get_flag()