/*
 * Copyright © 2016 Keith Packard <keithp@keithp.com>
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License along
 * with this program; if not, write to the Free Software Foundation, Inc.,
 * 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA.
 */

#include <stdio.h>
#include <errno.h>
#include <stdlib.h>
#include <unistd.h>
#include <libusb.h>
#include <getopt.h>
#include <string.h>
#include <strings.h>
#include "config.h"

#define CHAOS_SIZE	64

#define CHAOS_VENDOR	0x1d50
#define CHAOS_PRODUCT	0x60c6

struct chaoskey {
	libusb_context		*ctx;
	libusb_device_handle	*handle;
	int			kernel_active;
};

static libusb_device_handle *
chaoskey_match(libusb_device *dev, char *match_serial)
{
	struct libusb_device_descriptor desc;
	int ret;
	int match_len;
	char *device_serial = NULL;
	libusb_device_handle *handle = NULL;

	ret = libusb_get_device_descriptor(dev, &desc);
	if (ret < 0) {
		fprintf(stderr, "failed to get device descriptor: %s\n", libusb_strerror(ret));
		return 0;
	}

	if (desc.idVendor != CHAOS_VENDOR)
		return NULL;
	if (desc.idProduct != CHAOS_PRODUCT)
		return NULL;

	ret = libusb_open(dev, &handle);

	if (match_serial == NULL)
		return handle;

	if (ret < 0) {
		fprintf(stderr, "failed to open device: %s\n", libusb_strerror(ret));
		return NULL;
	}

	match_len = strlen(match_serial);
	device_serial = malloc(match_len + 2);

	if (!device_serial) {
		fprintf(stderr, "malloc failed\n");
		goto out;
	}

	ret = libusb_get_string_descriptor_ascii(handle, desc.iSerialNumber, (unsigned char *) device_serial, match_len + 1);

	if (ret < 0) {
		fprintf(stderr, "failed to get serial number: %s\n", libusb_strerror(ret));
		goto out;
	}

	device_serial[ret] = '\0';

	ret = strcmp(device_serial, match_serial);
	free(device_serial);
	if (ret)
		goto out;

	return handle;

out:
	if (handle)
		libusb_close(handle);
	return 0;
}

static struct chaoskey *
chaoskey_open(char *serial)
{
	struct chaoskey	*ck;
	int		ret;
	ssize_t		num;
	libusb_device	**list;
	int		d;

	ck = calloc(sizeof (struct chaoskey), 1);
	if (!ck)
		goto out;
	ret = libusb_init(&ck->ctx);
	if (ret ) {
		fprintf(stderr, "libusb_init failed: %s\n", libusb_strerror(ret));
		goto out;
	}

	num = libusb_get_device_list(ck->ctx, &list);
	if (num < 0) {
		fprintf(stderr, "libusb_get_device_list failed: %s\n", libusb_strerror(num));
		goto out;
	}

	for (d = 0; d < num; d++) {
		libusb_device_handle *handle;

		handle = chaoskey_match(list[d], serial);
		if (handle) {
			ck->handle = handle;
			break;
		}
	}

	libusb_free_device_list(list, 1);

	if (!ck->handle) {
		if (serial)
			fprintf (stderr, "No chaoskey matching %s\n", serial);
		else
			fprintf (stderr, "No chaoskey\n");
		goto out;
	}

	ck->kernel_active = libusb_kernel_driver_active(ck->handle, 0);
	if (ck->kernel_active) {
		ret = libusb_detach_kernel_driver(ck->handle, 0);
		if (ret)
			goto out;
	}

	ret = libusb_claim_interface(ck->handle, 0);
	if (ret)
		goto out;

	return ck;
out:
	if (ck->kernel_active)
		libusb_attach_kernel_driver(ck->handle, 0);
	if (ck->handle)
		libusb_close(ck->handle);
	if (ck->ctx)
		libusb_exit(ck->ctx);
	free(ck);
	return NULL;
}

static void
chaoskey_close(struct chaoskey *ck)
{
	libusb_release_interface(ck->handle, 0);
	if (ck->kernel_active)
		libusb_attach_kernel_driver(ck->handle, 0);
	libusb_close(ck->handle);
	libusb_exit(ck->ctx);
	free(ck);
}

#define COOKED_ENDPOINT	0x85
#define RAW_ENDPOINT	0x86
#define FLASH_ENDPOINT	0x87

static int
chaoskey_read(struct chaoskey *ck, int endpoint, void *buffer, int len)
{
	uint8_t	*buf = buffer;
	int	total = 0;

	while (len) {
		int	ret;
		int	transferred;

		ret = libusb_bulk_transfer(ck->handle, endpoint, buf, len, &transferred, 10000);
		if (ret) {
			if (total)
				return total;
			else {
				errno = EIO;
				return -1;
			}
		}
		len -= transferred;
		buf += transferred;
		total += transferred;
	}
	return total;
}

static const struct option options[] = {
	{ .name = "serial", .has_arg = 1, .val = 's' },
	{ .name = "length", .has_arg = 1, .val = 'l' },
	{ .name = "infinite", .has_arg = 0, .val = 'i' },
	{ .name = "bytes", .has_arg = 0, .val = 'b' },
	{ .name = "cooked", .has_arg = 0, .val = 'c' },
	{ .name = "raw", .has_arg = 0, .val = 'r' },
	{ .name = "flash", .has_arg = 0, .val = 'f' },
	{ .name = "version", .has_arg = 0, .val = 'V' },
	{ .name = "help", .has_arg = 0, .val = '?' },
	{ 0, 0, 0, 0},
};

static void usage(char *program, int code)
{
	fprintf(stderr, "usage: %s [--serial=<serial>] [--length=<length>[kMG]] [--infinite] [--bytes] [--cooked] [--raw] [--flash] [ --version] [--help]\n", program);
	exit(code);
}

static void
finish(struct chaoskey *ck, int code)
{
	chaoskey_close(ck);
	exit(code);
}

int
main (int argc, char **argv)
{
	struct chaoskey	*ck;
	uint16_t	buf[512];
	int	got;
	int	c;
	char	*serial = NULL;
	char	*length_string;
	char	*length_end;
	unsigned long	length = sizeof(buf);
	int	this_time;
	int	infinite = 0;
	int	bytes = 0;
	int	endpoint = RAW_ENDPOINT;

	while ((c = getopt_long(argc, argv, "s:l:ibcrfV?", options, NULL)) != -1) {
		switch (c) {
		case 's':
			serial = optarg;
			break;
		case 'l':
			length_string = optarg;
			length = strtoul(length_string, &length_end, 10);
			if (!strcasecmp(length_end, "k"))
				length *= 1024;
			else if (!strcasecmp(length_end, "m"))
				length *= 1024 * 1024;
			else if (!strcasecmp(length_end, "g"))
				length *= 1024 * 1024 * 1024;
			else if (strlen(length_end))
				usage(argv[0], 1);
			break;
		case 'i':
			infinite = 1;
			break;
		case 'b':
			bytes = 1;
			break;
		case 'c':
			endpoint = COOKED_ENDPOINT;
			break;
		case 'r':
			endpoint = RAW_ENDPOINT;
			break;
		case 'f':
			endpoint = FLASH_ENDPOINT;
			break;
		case 'V':
			printf("%s %s (%s)\n", PACKAGE, VERSION, RELEASE_DATE);
			exit(0);
		case '?':
			usage(argv[0], 0);
			break;
		default:
			usage(argv[0], 1);
			break;
		}
	}

	ck = chaoskey_open(serial);
	if (!ck)
		exit(1);

	if (bytes)
		length *= 2;

	while (length || infinite) {
		this_time = sizeof(buf);
		if (!infinite && length < sizeof(buf))
			this_time = (int) length;
		got = chaoskey_read(ck, endpoint, buf, this_time);
		if (got < 0) {
			perror("read");
			finish(ck, 1);
		}
		if (bytes) {
			int i;
			for (i = 0; i < got / 2; i++)
				putchar((buf[i] >> 1 & 0xff));
		} else {
			int i;
			int ret;
			for (i = 0; i < got; i += ret) {
				ret = write(1, ((char *) buf) + i, got - i);
				if (ret <= 0) {
					perror("write");
					finish(ck, 1);
				}
			}
		}
		length -= got;
	}
	finish(ck, 0);
}
