/*
 * hasplib.c: The guts of the Hasp HL USB library implementation.
 * 
 * Copyright (C) 2014 Jason A. Donenfeld <Jason@zx2c4.com>. All Rights Reserved.
 * 
 *  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 3 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, see <http://www.gnu.org/licenses/>.
 */

#include "hasplib.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <libusb.h>

struct hasp_response {
	uint8_t status;
	uint8_t status_checksum;
	uint8_t data[16];
};

struct _hasp_dongle {
	libusb_device_handle *handle;
	bool initialized;
	
	uint16_t lfsr;
	uint16_t polynomial_bottom;
	
	uint16_t password1;
	uint16_t password2;
};

enum HASP_FUNCTION_LIST {
	HASP_FUNCTION_SET_KEYS                 = 0x80,
	HASP_FUNCTION_CHECK_PASS               = 0x81,
	HASP_FUNCTION_READ_3WORDS              = 0x82,
	HASP_FUNCTION_WRITE_WORD               = 0x83,
	HASP_FUNCTION_READ_ST                  = 0x84,
	HASP_FUNCTION_READ_NETMEMORY_3WORDS    = 0x8B,
	HASP_FUNCTION_HASH_DWORD               = 0x98
};

enum HASP_STATUS {
	HASP_STATUS_OK                     = 0x00,
	HASP_STATUS_ERROR                  = 0x01,
	HASP_STATUS_INVALID_MEMORY_ADDRESS = 0x04,
	HASP_STATUS_LAST                   = 0x1F,
	HASP_USB_ERROR                     = 0x6D
};

static void cipher(hasp_dongle *dongle, uint8_t *buffer, size_t len)
{
	const uint16_t polynomial = (dongle->polynomial_bottom >> 1) | 0x8000;
	uint16_t lfsr = dongle->lfsr;
	uint8_t stream_byte, output_bit;

	while (len--) {
		stream_byte = 0;
		for (int i = 0; i < 4; ++i) {
			/* Output a bit from position 0 into stream_byte. */
			output_bit = lfsr & 1;
			stream_byte = (stream_byte << 1) | output_bit;
			lfsr >>= 1;

			/* Flip the tap bits if we output a 1. */
			lfsr ^= (-output_bit) & polynomial;

			/* Output a bit from position 7 into stream_byte. */
			output_bit = (lfsr >> 7) & 1;
			stream_byte = (stream_byte << 1) | output_bit;
		}
		*(buffer++) ^= stream_byte;
	}
	dongle->lfsr = lfsr;
}

static void usb_error(int error)
{
	fprintf(stderr, "USB error: %s\n", libusb_strerror(error));
}

static void make_request(hasp_dongle *dongle, uint8_t function, uint16_t param1, uint16_t param2, size_t expected_data_len, int cipher_params, bool retry, struct hasp_response *response)
{
	int ret;
	response->status = HASP_STATUS_ERROR;
	memset(response->data, 0, sizeof(response->data));
	
	uint16_t params[2] = { param1, param2 };
	if (cipher_params)
		cipher(dongle, (uint8_t *)params, cipher_params * sizeof(uint16_t));

	ret = libusb_control_transfer(dongle->handle, 0xC0, function, params[0], params[1], (unsigned char *)response, expected_data_len + 2, 3000);
	if (ret < 0) {
		usb_error(ret);
		response->status = HASP_USB_ERROR;
		return;
	}
	
	cipher(dongle, (uint8_t *)response, expected_data_len + 2);
	
	if (response->status == HASP_STATUS_OK)
		   dongle->polynomial_bottom = (dongle->polynomial_bottom & 0xFF) | (response->status_checksum << 8);
	else if (retry && dongle->password1 && dongle->password2) {
		if (hasp_init(dongle) && hasp_login(dongle, dongle->password1, dongle->password2, NULL))
			make_request(dongle, function, param1, param2, expected_data_len, cipher_params, false, response);
	}
}

size_t hasp_find_dongles(hasp_dongle ***dongles)
{
	size_t dongle_count;
	ssize_t size;
	libusb_device **all_devs, **all_pos, *dev;
	hasp_dongle **pos, *dongle;
	int ret;

	if ((ret = libusb_init(NULL)) < 0) {
		usb_error(ret);
		return 0;
	}

	if ((size = libusb_get_device_list(NULL, &all_devs)) < 0) {
		usb_error((int)size);
		return 0;
	}
	if (!size)
		return 0;
	all_pos = all_devs;

	pos = calloc(size + 1, sizeof(hasp_dongle *));
	if (!pos)
		return 0;
	*dongles = pos;

	dongle_count = 0;
	while ((dev = *(all_pos++)) != NULL) {
		struct libusb_device_descriptor desc;
		ret = libusb_get_device_descriptor(dev, &desc);
		if (ret < 0) {
			usb_error(ret);
			continue;
		}
		if (desc.idVendor == 0x0529 && desc.idProduct == 0x0001) {
			libusb_device_handle *handle;
			if ((ret = libusb_open(dev, &handle)) < 0) {
				usb_error(ret);
				continue;
			}
			
			dongle = calloc(1, sizeof(hasp_dongle));
			if (!dongle)
				return 0;
			dongle->handle = handle;

			++dongle_count;
			*(pos++) = dongle;
		}
	}
	libusb_free_device_list(all_devs, 1);
	return dongle_count;
}
hasp_dongle *hasp_find_login_first_dongle(uint16_t password1, uint16_t password2)
{
	hasp_dongle *found = NULL;
	hasp_dongle **all = NULL;
	size_t num_dongles;
	if (!(num_dongles = hasp_find_dongles(&all))) {
		free(all);
		return NULL;
	}
	for (size_t i = 0; i < num_dongles; ++i) {
		if (!found && hasp_login(all[i], password1, password2, NULL))
			found = all[i];
		else
			hasp_free_dongle(all[i]);
	}
	free(all);
	return found;
}

void hasp_free_dongles(hasp_dongle **dongles)
{
	for (hasp_dongle **i = dongles; *i; ++i)
		hasp_free_dongle(*i);
	free(dongles);
}

void hasp_free_dongle(hasp_dongle *dongle)
{
	if (dongle->handle)
		libusb_close(dongle->handle);
	free(dongle);
}

bool hasp_init(hasp_dongle *dongle)
{
	dongle->lfsr = 0xFABE;
	dongle->polynomial_bottom = 0xA0CB;
	
	struct hasp_response response;
	make_request(dongle, HASP_FUNCTION_SET_KEYS, dongle->lfsr, 0 /* unused */, 5, false, false, &response);
	
	if (response.status == HASP_STATUS_OK && response.data[0] == 0x2 && response.data[1] == 0xEA /* 0x0A? */ && response.data[2] == 0x00) {
		dongle->initialized = true;
		return true;
	}
	return false;
}

bool hasp_login(hasp_dongle *dongle, uint16_t password1, uint16_t password2, uint16_t *memory_size)
{
	dongle->password1 = password1;
	dongle->password2 = password2;
	
	if (!dongle->initialized && !hasp_init(dongle))
		return false;
	
	struct hasp_response response;
	make_request(dongle, HASP_FUNCTION_CHECK_PASS, dongle->password1, dongle->password2, 3, 2, false, &response);
	if (response.status != HASP_STATUS_OK && response.status != HASP_USB_ERROR) {
		hasp_init(dongle);
		make_request(dongle, HASP_FUNCTION_CHECK_PASS, dongle->password1, dongle->password2, 3, 2, false, &response);
	}
	
	if (response.status != HASP_STATUS_OK && response.data[2] != 0x10)
		return false;
	if (memory_size)
		*memory_size = response.data[0] | (response.data[1] << 8);
	return true;
}

bool hasp_id(hasp_dongle *dongle, uint32_t *id)
{
	struct hasp_response response;
	make_request(dongle, HASP_FUNCTION_READ_NETMEMORY_3WORDS, 0 /* first location has id */, 0 /* unused */, 6, 1, true, &response);
	if (response.status != HASP_STATUS_OK)
		return false;
	*id = response.data[0] | (response.data[1] << 8) | (response.data[2] << 16) | (response.data[3] << 24);
	return true;
}

bool hasp_read(hasp_dongle *dongle, uint16_t location, uint16_t *value)
{
	struct hasp_response response;
	make_request(dongle, HASP_FUNCTION_READ_3WORDS, location, 0 /* unused */, 6, 1, true, &response);
	if (response.status != HASP_STATUS_OK)
		return false;
	*value = response.data[0] | (response.data[1] << 8);
	return true;
}

bool hasp_write(hasp_dongle *dongle, uint16_t location, uint16_t value)
{
	struct hasp_response response;
	make_request(dongle, HASP_FUNCTION_WRITE_WORD, location, value, 0, 2, true, &response);
	return response.status == HASP_STATUS_OK;
}