#include <stdio.h>
#include <stdlib.h>
#include <stdarg.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <string.h>
#include <libusb-1.0/libusb.h>

#include "config.h"

#define ERR_EXIT(errcode) do { perr("   %s\n", libusb_strerror((enum libusb_error)errcode)); return -1; } while (0)
#define CALL_CHECK(fcall) do { int _r=fcall; if (_r < 0) ERR_EXIT(_r); } while (0)
#define CALL_CHECK_CLOSE(fcall, hdl) do { int _r=fcall; if (_r < 0) { libusb_close(hdl); ERR_EXIT(_r); } } while (0)

#define REQUEST_SENSE_LENGTH 0x12

// Section 5.1: Command Block Wrapper (CBW)
struct command_block_wrapper {
    uint8_t dCBWSignature[4];
    uint32_t dCBWTag;
    uint32_t dCBWDataTransferLength;
    uint8_t bmCBWFlags;
    uint8_t bCBWLUN;
    uint8_t bCBWCBLength;
    uint8_t CBWCB[16];
};

// Section 5.2: Command Status Wrapper (CSW)
struct command_status_wrapper {
    uint8_t dCSWSignature[4];
    uint32_t dCSWTag;
    uint32_t dCSWDataResidue;
    uint8_t bCSWStatus;
};

static const uint8_t cdb_length[256] = {
  // 0  1  2  3  4  5  6  7  8  9  A  B  C  D  E  F
    06,06,06,06,06,06,06,06,06,06,06,06,06,06,06,06,  //  0
    06,06,06,06,06,06,06,06,06,06,06,06,06,06,06,06,  //  1
    10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,  //  2
    10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,  //  3
    10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,  //  4
    10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,10,  //  5
    00,00,00,00,00,00,00,00,00,00,00,00,00,00,00,00,  //  6
    00,00,00,00,00,00,00,00,00,00,00,00,00,00,00,00,  //  7
    16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,  //  8
    16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,16,  //  9
    12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,  //  A
    12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,12,  //  B
    00,00,00,00,00,00,00,00,00,00,00,00,00,00,00,00,  //  C
    00,00,00,00,00,00,00,00,00,00,00,00,00,00,00,00,  //  D
    00,00,00,00,00,00,00,00,00,00,00,00,00,00,00,00,  //  E
    00,00,00,00,00,00,00,00,00,00,00,00,00,16,16,16,  //  F
};

// static struct libusb_device_handle *devh = NULL;

static void perr(char const *format, ...)
{
	va_list args;

	va_start (args, format);
	vfprintf(stderr, format, args);
	va_end(args);
}

static void display_buffer_hex(unsigned char *buffer, unsigned size)
{
	unsigned i, j, k;

	for (i=0; i<size; i+=16) {
		printf("\n  %08x  ", i);
		for(j=0,k=0; k<16; j++,k++) {
			if (i+j < size) {
				printf("%02x", buffer[i+j]);
			} else {
				printf("  ");
			}
			printf(" ");
		}
		printf(" ");
		for(j=0,k=0; k<16; j++,k++) {
			if (i+j < size) {
				if ((buffer[i+j] < 32) || (buffer[i+j] > 126)) {
					printf(".");
				} else {
					printf("%c", buffer[i+j]);
				}
			}
		}
	}
	printf("\n" );
}

static int send_mass_storage_command(libusb_device_handle *handle, uint8_t endpoint, uint8_t lun,
        uint8_t *cdb, uint8_t direction, int data_length, uint32_t *ret_tag)
{
    static uint32_t tag = 1;
    uint8_t cdb_len;
    int i, r, size;
    struct command_block_wrapper cbw;

    if (cdb == NULL) {
        return -1;
    }

    if (endpoint & LIBUSB_ENDPOINT_IN) {
        perr("send_mass_storage_command: cannot send command on IN endpoint\n");
        return -1;
    }

    display_buffer_hex(cdb, cdb_length[cdb[0]]);

    cdb_len = cdb_length[cdb[0]];
    if ((cdb_len == 0) || (cdb_len > sizeof(cbw.CBWCB))) {
        perr("send_mass_storage_command: don't know how to handle this command (%02X, length %d)\n",
                cdb[0], cdb_len);
        return -1;
    }

    memset(&cbw, 0, sizeof(cbw));
    cbw.dCBWSignature[0] = 'U';
    cbw.dCBWSignature[1] = 'S';
    cbw.dCBWSignature[2] = 'B';
    cbw.dCBWSignature[3] = 'C';
    *ret_tag = tag;
    cbw.dCBWTag = tag++;
    cbw.dCBWDataTransferLength = data_length;
    cbw.bmCBWFlags = direction;
    cbw.bCBWLUN = lun;
    // Subclass is 1 or 6 => cdb_len
    cbw.bCBWCBLength = cdb_len;
    memcpy(cbw.CBWCB, cdb, cdb_len);

    i = 0;
    do {
        // The transfer length must always be exactly 31 bytes.
        r = libusb_bulk_transfer(handle, endpoint, (unsigned char*)&cbw, 31, &size, 1000);
        if (r == LIBUSB_ERROR_PIPE) {
            libusb_clear_halt(handle, endpoint);
        }
        i++;
    } while ((r == LIBUSB_ERROR_PIPE) && (i<RETRY_MAX));
    if (r != LIBUSB_SUCCESS) {
        perr("send_mass_storage_command: %s\n", libusb_strerror((enum libusb_error)r));
        return -1;
    }

    printf("sent %d CDB bytes\n", cdb_len);
    return 0;
}

static int get_mass_storage_status(libusb_device_handle *handle, uint8_t endpoint, uint32_t expected_tag)
{
    int i, r, size;
    struct command_status_wrapper csw;

    // The device is allowed to STALL this transfer. If it does, you have to
    // clear the stall and try again.
    i = 0;
    do {
        r = libusb_bulk_transfer(handle, endpoint, (unsigned char*)&csw, 13, &size, 1000);
        if (r == LIBUSB_ERROR_PIPE) {
            libusb_clear_halt(handle, endpoint);
        }
        i++;
    } while ((r == LIBUSB_ERROR_PIPE) && (i<RETRY_MAX));
    if (r != LIBUSB_SUCCESS) {
        perr("get_mass_storage_status: %s\n", libusb_strerror((enum libusb_error)r));
        return -1;
    }
    if (size != 13) {
        perr("get_mass_storage_status: received %d bytes (expected 13)\n", size);
        return -1;
    }
    if (csw.dCSWTag != expected_tag) {
        perr("get_mass_storage_status: mismatched tags (expected %08X, received %08X)\n",
                expected_tag, csw.dCSWTag);
        return -1;
    }
    // For this test, we ignore the dCSWSignature check for validity...
    printf("Mass Storage Status: %02X (%s)\n", csw.bCSWStatus, csw.bCSWStatus?"FAILED":"Success");
    if (csw.dCSWTag != expected_tag)
        return -1;
    if (csw.bCSWStatus) {
        // REQUEST SENSE is appropriate only if bCSWStatus is 1, meaning that the
        // command failed somehow.  Larger values (2 in particular) mean that
        // the command couldn't be understood.
        if (csw.bCSWStatus == 1)
            return -2;	// request Get Sense
        else
            return -1;
    }

    // In theory we also should check dCSWDataResidue.  But lots of devices
    // set it wrongly.
    return 0;
}

int main(int argc, char *argv[])
{
    int r = 0, e = 0;
    struct libusb_device_handle *handle = NULL;
    struct libusb_device **devs;
    struct libusb_device *dev;
    struct libusb_device_descriptor desc;
    
    /* Init libusb */
    r = libusb_init(NULL);
    if (r < 0)
    {
        printf("\nfailed to initialise libusb\n");
        return 1;
    }

    handle = libusb_open_device_with_vid_pid(NULL, VENDOR_ID, PRODUCT_ID);
    if(handle == NULL) {
        printf("\nError in device opening!");
    } else {
        printf("\nDevice Opened");
    }

    libusb_set_configuration(handle, 1);
    if(libusb_kernel_driver_active(handle, 0) == 1) {
        printf("\nKernel Driver Active");
        if(libusb_detach_kernel_driver(handle, 0) == 0) {
            printf("\nKernel Driver Detached!");
        }
    }

    e = libusb_claim_interface(handle, 0);
    if(e < 0) {
        printf("\nCannot Claim Interface"); 
    } else {
        printf("\nClaimed Interface");
    }

    uint32_t expected_tag = 0;
    uint8_t cdb[16];
    uint8_t sense[18];
    int size;
    int rc;

    /* start: set passphrase */
    memset(cdb, 0, sizeof(cdb));
    cdb[0] = 0x03;
    cdb[4] = REQUEST_SENSE_LENGTH;
    send_mass_storage_command(handle, 0x01, 0, cdb, LIBUSB_ENDPOINT_IN, REQUEST_SENSE_LENGTH, &expected_tag);
    rc = libusb_bulk_transfer(handle, 0x81, (unsigned char*)&sense, REQUEST_SENSE_LENGTH, &size, 1000);
    if (rc < 0)
    {
        printf("libusb_bulk_transfer failed: %s\n", libusb_error_name(rc));
        return 1;
    }
    printf("received %d bytes\n", size);

    rc = libusb_bulk_transfer(handle, 0x81, (unsigned char*)&sense, REQUEST_SENSE_LENGTH, &size, 1000);
    if (rc < 0)
    {
        printf("libusb_bulk_transfer failed: %s\n", libusb_error_name(rc));
        return 1;
    }
    printf("received %d bytes\n", size);

    memset(cdb, 0, sizeof(cdb));
    cdb[0] = 0x03;
    cdb[4] = REQUEST_SENSE_LENGTH;

    unsigned char cmd[] = "\xfe\x00\x37\x37\x37\x37\x37\x37\x37\x37\x37\x37\x00\x00\x00\x00";
    send_mass_storage_command(handle, 0x01, 0, cmd, LIBUSB_ENDPOINT_IN, REQUEST_SENSE_LENGTH, &expected_tag);
    rc = libusb_bulk_transfer(handle, 0x81, (unsigned char*)&sense, REQUEST_SENSE_LENGTH, &size, 1000);
    if (rc < 0)
    {
        printf("libusb_bulk_transfer failed: %s\n", libusb_error_name(rc));
        return 1;
    }
    printf("received %d bytes\n", size);
    /* end: set passphrase */

    /* start: set ssid */
    memset(cdb, 0, sizeof(cdb));
    cdb[0] = 0x03;
    cdb[4] = REQUEST_SENSE_LENGTH;
    send_mass_storage_command(handle, 0x01, 0, cdb, LIBUSB_ENDPOINT_IN, REQUEST_SENSE_LENGTH, &expected_tag);
    rc = libusb_bulk_transfer(handle, 0x81, (unsigned char*)&sense, REQUEST_SENSE_LENGTH, &size, 1000);
    if (rc < 0)
    {
        printf("libusb_bulk_transfer failed: %s\n", libusb_error_name(rc));
        return 1;
    }
    printf("received %d bytes\n", size);

    rc = libusb_bulk_transfer(handle, 0x81, (unsigned char*)&sense, REQUEST_SENSE_LENGTH, &size, 1000);
    if (rc < 0)
    {
        printf("libusb_bulk_transfer failed: %s\n", libusb_error_name(rc));
        return 1;
    }
    printf("received %d bytes\n", size);

    /* set ssid */
    memset(cdb, 0, sizeof(cdb));
    cdb[0] = 0x03;
    cdb[4] = REQUEST_SENSE_LENGTH;

    unsigned char cmd2[] = "\xfd\x00\x38\x38\x38\x38\x38\x38\x38\x38\x38\x38\x00\x00\x00\x00";
    send_mass_storage_command(handle, 0x01, 0, cmd2, LIBUSB_ENDPOINT_IN, REQUEST_SENSE_LENGTH, &expected_tag);
    rc = libusb_bulk_transfer(handle, 0x81, (unsigned char*)&sense, REQUEST_SENSE_LENGTH, &size, 1000);
    if (rc < 0)
    {
        printf("libusb_bulk_transfer failed: %s\n", libusb_error_name(rc));
        return 1;
    }
    printf("received %d bytes\n", size);
    /* end: set ssid */

    //get_mass_storage_status(handle, 0x01, expected_tag);

    libusb_release_interface(handle, 0);
    libusb_reset_device(handle);
    libusb_close(handle);

    libusb_exit(NULL);

    exit(EXIT_SUCCESS);
}