/*
 * Kangaroo Punch MultiPlayer Game Server Mark II
 * Copyright (C) 2020-2021  Scott Duensing
 *
 * 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 <https://www.gnu.org/licenses/>.
 *
 */


// Reliability:   https://github.com/BaroboRobotics/libsfp/wiki/Serial-Framing-Protocol
// Key Exchange:  https://www.techiedelight.com/c-program-demonstrate-diffie-hellman-algorithm
// Encryption:    https://erev0s.com/blog/tiny-aes-cbc-mode-pkcs7-padding-written-c


/*
 * NOTE:
 *
 * The reliable transport / packet retry code has been commented out.
 * Structly speaking, it isn't needed.  Enet handles reliable delivery
 * across the internet.  The only place there can be corruption is
 * between a real serial port and the enet softmodem.
 *
 * In any case, if the code ends up being re-enabled, the sequence
 * numbers have no bounding and will run off the end of the history
 * array.  Not only that, but due to the limited number of sequence
 * numbers available, there needs to be a sending queue added for
 * when the server attempts to send more packets than there are
 * history slots.
 *
 */


#include <stdarg.h>

#include "packet.h"
#include "primes.h"

#define CBC 1
#include "thirdparty/tiny-AES-c/aes.h"
#include "thirdparty/tiny-AES128-C/pkcs7_padding.h"


static packetSender _packetSender    = NULL;
static uint8_t      _encryptionReady = 0;


static void     *packetAesContextCreate(uint8_t *key, uint8_t *iv);
static uint8_t   packetCRC(char *data, uint16_t length, uint8_t startAt);
static uint16_t  packetDHCompute(uint32_t a, uint32_t m, uint32_t n);


static void *packetAesContextCreate(uint8_t *key, uint8_t *iv) {
	struct AES_ctx *ctx         = NULL;

	NEW(struct AES_ctx, ctx);
	if (ctx) {
		// Both bits parameters are 128 bits in length (16 bytes).
		AES_init_ctx_iv(ctx, key, iv);
	}

	return ctx;
}


char *packetContentPack(uint16_t *length, char *format, ...) {
	va_list      args             = { 0 };
	static char  work[PACKET_MAX] = { 0 };
	char        *result           = NULL;
	char        *buffer           = NULL;
	int32_t      value32          = 0;

	*length = 0;

	va_start(args, format);
	if (format) {
		while (*format != 0) {
			switch (*format) {
				case 'i':
					// Copy integer as 32-bit integer.
					value32 = va_arg(args, int32_t);
					memcpy(&work[*length], &value32, sizeof(int32_t));
					*length += sizeof(int32_t);
					break;

				case 's':
					// Copy string including the terminating zero.
					buffer = va_arg(args, char *);
					memcpy(&work[*length], buffer, strlen(buffer) + 1);
					*length += strlen(buffer) + 1;
					break;

				default:
					utilDie("restRequest:  Unknown format option '%c'.\n", *format);
					break;
			}
			format++;
		}
	}
	va_end(args);

	// Copy working buffer to result to return.
	if (*length > 0) {
		result = (char *)malloc(*length);
		memcpy(result, work, *length);
	}

	return result;
}


void packetContentUnpack(char *buffer, char *format, ...) {
	va_list      args     = { 0 };
	char       **string   = NULL;
	int32_t     *value32  = NULL;
	uint16_t     length   = 0;

	va_start(args, format);
	if (format) {
		while (*format != 0) {
			switch (*format) {
				case 'i':
					// Copy integer as 32-bit integer.
					value32 = va_arg(args, int32_t *);
					memcpy(value32, &buffer[length], sizeof(int32_t));
					length += sizeof(int32_t);
					break;

				case 's':
					// Copy string including the terminating zero.
					string = va_arg(args, char **);
					if (&buffer[length] != NULL) {
						*string = strdup(&buffer[length]);
						length += strlen(*string) + 1;
					}
					break;

				default:
					utilDie("restRequest:  Unknown format option '%c'.\n", *format);
					break;
			}
			format++;
		}
	}
	va_end(args);
}


static uint8_t packetCRC(char *data, uint16_t length, uint8_t startAt) {
	uint16_t x = 0;

	/*
	 * ***TODO***
	 *
	 * Something in the data is throwing this off.  CRCs fail when there are
	 * "special characters" in the data, such as these string table fields:
	 *
	 * Storing [userAllowed]=[abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890`!@#$%^&*()_-+=[]{}\/?<>.,:;"']
	 * CRC failed!  192 != 60
	 *
	 * Storing [passSpecial]=[`!@#$%^&*()_-+=[]\{}|:;"',.<>/?]
	 * CRC failed!  8 != 244
	 *
	 * Like the rest of the reliability code, this is only needed if we have
	 * RS-232 issues to the softmodem.  Documenting for later if needed.
	 *
	 */

	for (x=0; x<length; x++) {
		startAt ^= data[x] + data[x];  // Just XOR and increment.
	}

	return startAt;
}


void packetDebugPrint(uint8_t *data, uint16_t length) {
	uint16_t  i;
	for (i=0; i<length; i++) logWrite("%02x [%c] ", data[i], data[i]);
	logWrite("\n\r");
}


uint8_t packetDecode(PacketThreadDataT *threadData, PacketDecodeDataT *decodeData, char *input, uint16_t inputLength) {
	//uint8_t           sequence = 0;
	uint16_t          unpadded = 0;
	int32_t           x        = 0;
	uint8_t           c        = 0;
	PacketEncodeDataT encoded  = { 0 };

	// input and inputLength are incoming raw data or NULL and 0 to continue processing already received data.
	// Returns 1 on packet ready, 0 on still waiting.

	// Is there input data to add to the queue?
	if (input != NULL && inputLength > 0) {
		for (x=0; x<inputLength; x++) {
			threadData->decodeQueue[threadData->decodeQueueHead++] = input[x];
			if (threadData->decodeQueueHead >= PACKET_INPUT_QUEUE_SIZE) threadData->decodeQueueHead = 0;
		}
	}

	while (1) {
		// Do we have data to process?
		if (threadData->decodeQueueHead == threadData->decodeQueueTail) return 0;

		// Get next byte.
		c = threadData->decodeQueue[threadData->decodeQueueTail++];
		if (threadData->decodeQueueTail >= PACKET_INPUT_QUEUE_SIZE) threadData->decodeQueueTail = 0;

		// New packet?
		if (threadData->newPacket) {
			threadData->newPacket = 0;
			threadData->inEscape  = 0;
			threadData->length    = 0;
		}

		// Are we escaped?
		if (threadData->inEscape) {
			threadData->inEscape = 0;
			// Is this the end of the packet?
			if (c == PACKET_FRAME) {
				// Nope!  Add the byte.
				threadData->decodeBuffer[threadData->length++] = c;
			} else {
				// Yes!  Process it and set up for the next packet.
				threadData->newPacket = 1;

				/*
				// Check CRC.
				x = packetCRC(threadData->decodeBuffer, threadData->length - 1, 0);
				if ((uint8_t)threadData->decodeBuffer[threadData->length - 1] != (uint8_t)x) {
					//logWrite("CRC failed!  %d != %d\n", (uint8_t)threadData->decodeBuffer[threadData->length - 1], (uint8_t)x);
					continue;
				}

				// Get sequence value.
				sequence = ((uint8_t)threadData->decodeBuffer[0]) & 0x1f;

				// Is this a NAK?
				if ((((uint8_t)threadData->decodeBuffer[0]) & 0xc0) >> 6 == PACKET_CONTROL_NAK) {
					// Rewind until we find the packet we need in history.
					x = threadData->historyPosition - 1;
					if (x < 0) x = PACKET_SEQUENCE_MAX - 1;
					while (threadData->history[x].sequence != sequence && x != threadData->historyPosition) {
						x--;
						if (x < 0) x = PACKET_SEQUENCE_MAX - 1;
					}
					// Did we find it?
					if (x == threadData->historyPosition) {
						// No!  BAD!
						logWrite("Unable to locate missing packet in history!\n\r");
					} else {
						// Yes.  Replay missing packets.
						while (x != threadData->historyPosition) {
							logWrite("Resending %d!\n\r", threadData->history[x].sequence);
							_packetSender(threadData->history[x].data, threadData->history[x].length, threadData->senderData);
							x++;
							if (x >= PACKET_SEQUENCE_MAX) x = 0;
						}
					}
					continue;
				}

				// Is this the sequence number we're expecting?
				if (sequence == threadData->lastRemoteSequence) {
					// Yes!
					threadData->lastRemoteSequence++;
				} else {
					// No!  NAK it!
					logWrite("Packet out of sequence!  Got %d wanted %d!\n\r", sequence, threadData->lastRemoteSequence);
					encoded.control    = PACKET_CONTROL_NAK;              // Negative acknowledge.
					encoded.packetType = PACKET_TYPE_NONE;                // Not destined for the app.
					encoded.channel    = 0;                               // Channel doesn't matter for NAK.
					encoded.encrypt    = 0;                               // Encryption doesn't matter for NAK.
					encoded.sequence   = threadData->lastRemoteSequence;  // The last good packet we saw.
					packetEncode(threadData, &encoded, NULL, 0);
					packetSend(threadData, &encoded);
					continue;
				}
				*/

				// Fill decoded data fields.
				decodeData->packetType = threadData->decodeBuffer[1];
				decodeData->channel    = threadData->decodeBuffer[2];

				// AES Decryption.
				if (threadData->decodeBuffer[0] & 32) {
					AES_ctx_set_iv(threadData->aesContext, (uint8_t *)threadData->dhModulus);
					AES_CBC_decrypt_buffer(threadData->aesContext, (uint8_t *)&threadData->decodeBuffer[3], threadData->length - 4);
					unpadded = pkcs7_padding_data_length((uint8_t *)&threadData->decodeBuffer[3], threadData->length - 4, 16);
					// Fix length.
					threadData->length = unpadded + 4;
				}

				// Copy packet data to new buffer, if any.
				if (threadData->length - 4 > 0) {
					decodeData->data = (char *)malloc(threadData->length - 4);  // 4 for 3 byte header and 1 byte CRC.
					if (!decodeData) continue;
					memcpy(decodeData->data, &threadData->decodeBuffer[3], threadData->length - 4);  // Skip header and CRC.
					decodeData->length = threadData->length - 4;
				} else {
					// No payload.
					decodeData->data = NULL;
					decodeData->length = 0;
				}
				// Fix length to remove header and checksum.
				threadData->length -= 4;

				// Is this a DH_REQUEST?
				if (decodeData->packetType == PACKET_TYPE_DH_REQUEST) {
					memcpy(threadData->dhModulus,      decodeData->data,                              PACKET_ENCRYPT_KEY_SIZE * sizeof(uint16_t));
					memcpy(threadData->dhBase,        &decodeData->data[PACKET_ENCRYPT_KEY_SIZE * 2], PACKET_ENCRYPT_KEY_SIZE * sizeof(uint16_t));
					memcpy(threadData->dhTheirPublic, &decodeData->data[PACKET_ENCRYPT_KEY_SIZE * 4], PACKET_ENCRYPT_KEY_SIZE * sizeof(uint16_t));
					DEL(decodeData->data);
					for (x=0; x<PACKET_ENCRYPT_KEY_SIZE; x++) {
						threadData->dhMySecret[x]  = rand();
						threadData->dhMyPublic[x]  = packetDHCompute(threadData->dhBase[x],        threadData->dhMySecret[x], threadData->dhModulus[x]);
						threadData->dhSharedKey[x] = packetDHCompute(threadData->dhTheirPublic[x], threadData->dhMySecret[x], threadData->dhModulus[x]);
					}
					threadData->aesContext = packetAesContextCreate((uint8_t *)threadData->dhSharedKey, (uint8_t *)threadData->dhModulus);
					encoded.control    = PACKET_CONTROL_DAT;
					encoded.packetType = PACKET_TYPE_DH_RESPONSE;
					encoded.channel    = 0;  // Doesn't matter for DH_RESPONSE.
					encoded.encrypt    = 0;  // Must be 0 for DH_RESPONSE.
					packetEncode(threadData, &encoded, (char *)threadData->dhMyPublic, PACKET_ENCRYPT_KEY_SIZE * sizeof(uint16_t));
					packetSend(threadData, &encoded);
					_encryptionReady = 1;  // Responding side.
					continue;
				}

				// Is this a DH_RESPONSE?
				if (decodeData->packetType == PACKET_TYPE_DH_RESPONSE) {
					memcpy(threadData->dhTheirPublic, decodeData->data, PACKET_ENCRYPT_KEY_SIZE * sizeof(uint16_t));
					DEL(decodeData->data);
					for (x=0; x<PACKET_ENCRYPT_KEY_SIZE; x++) {
						threadData->dhSharedKey[x] = packetDHCompute(threadData->dhTheirPublic[x], threadData->dhMySecret[x], threadData->dhModulus[x]);
					}
					threadData->aesContext = packetAesContextCreate((uint8_t *)threadData->dhSharedKey, (uint8_t *)threadData->dhModulus);
					_encryptionReady = 1;  // Initiating side.
					continue;
				}

				// Done!
				break;
			}
		} else {
			// Are we escaping?
			if (c == PACKET_FRAME) {
				// Yes.  Don't add this byte.
				threadData->inEscape = 1;
				continue;
			}
			// Did we overflow?
			if (threadData->length >= PACKET_MAX) {
				// Yup.  Dump it.
				threadData->newPacket = 1;
				continue;
			}
			// Add byte to packet.
			threadData->decodeBuffer[threadData->length++] = c;
		}
	}

	return 1;
}


void packetDecodeDataDestroy(PacketDecodeDataT **packet) {
	PacketDecodeDataT *d = *packet;

	packetDecodeDataStaticDestroy(d);
	free(d);
	d = NULL;
	*packet = d;
}


void packetDecodeDataStaticDestroy(PacketDecodeDataT *packet) {
	if (packet->data) {
		free(packet->data);
		packet->data = NULL;
	}
}


static uint16_t packetDHCompute(uint32_t a, uint32_t m, uint32_t n) {
	uint32_t r = 0;
	uint32_t y = 1;

	while (m > 0) {
		r = m % 2;
		// Fast exponention.
		if (r == 1) {
			y = (y * a) % n;
		}
		a = a * a % n;
		m = m / 2;
	}

	return y;
}


uint8_t packetEncode(PacketThreadDataT *threadData, PacketEncodeDataT *data, char *input, uint16_t length) {
	uint16_t  x           = 0;
	uint8_t   crc         = 0;
	uint8_t   control     = 0;
	uint16_t  paddedSize  = 0;
	uint16_t  inputLength = 0;
	char     *inputBuffer = NULL;

	// Returns 1 on success, 0 on failure.

	// Packet too large?
	if (length > PACKET_MAX) return 0;

	data->dataPointer = NULL;
	data->length      = 0;

	/*
	if (data->control == PACKET_CONTROL_DAT) {
		data->sequence    = threadData->sequence++;
	}
	*/

	// Make needed header bytes.
	control = (((uint8_t)data->control) << 6) + (data->encrypt << 5) + (data->sequence & 0x1f);

	// Calculate CRC over header bytes.
	crc = packetCRC((char *)&control,          1, crc);
	crc = packetCRC((char *)&data->packetType, 1, crc);
	crc = packetCRC((char *)&data->channel,    1, crc);

	// Add header bytes.
	threadData->encodeBuffer[data->length++] = control;
	if (control == PACKET_FRAME) threadData->encodeBuffer[data->length++] = PACKET_FRAME;
	threadData->encodeBuffer[data->length++] = data->packetType;
	if (data->packetType == PACKET_FRAME) threadData->encodeBuffer[data->length++] = PACKET_FRAME;
	threadData->encodeBuffer[data->length++] = data->channel;
	if (data->channel == PACKET_FRAME) threadData->encodeBuffer[data->length++] = PACKET_FRAME;

	// AES Encryption.
	if (data->encrypt) {
		AES_ctx_set_iv(threadData->aesContext, (uint8_t *)threadData->dhModulus);
		if (length % 16) {
			paddedSize = length + 16 - (length % 16);
		} else {
			paddedSize = length;
		}
		memset(threadData->encryptionBuffer, 0, PACKET_MAX);
		for (x=0; x<length; x++) threadData->encryptionBuffer[x] = input[x];
		pkcs7_padding_pad_buffer(threadData->encryptionBuffer, length, paddedSize, 16);
		AES_CBC_encrypt_buffer(threadData->aesContext, threadData->encryptionBuffer, paddedSize);
		// Encrypted, select proper buffer.
		inputBuffer = (char *)threadData->encryptionBuffer;
		inputLength = paddedSize;
	} else {
		// Not encrypted, select proper buffer.
		inputBuffer = input;
		inputLength = length;
	}

	// Add payload.
	for (x=0; x<inputLength; x++) {
		// Is this a frame character?  If so, escape it.
		if ((uint8_t)inputBuffer[x] == PACKET_FRAME) threadData->encodeBuffer[data->length++] = PACKET_FRAME;
		// CRC data.
		crc = packetCRC(&inputBuffer[x], 1, crc);
		// Add data.
		threadData->encodeBuffer[data->length++] = inputBuffer[x];
	}

	// Add CRC.
	threadData->encodeBuffer[data->length++] = crc;
	if (crc == PACKET_FRAME) threadData->encodeBuffer[data->length++] = PACKET_FRAME;

	// Mark end of packet.
	threadData->encodeBuffer[data->length++] = PACKET_FRAME;
	threadData->encodeBuffer[data->length++] = 0;

	data->dataPointer = threadData->encodeBuffer;

	return 1;
}


uint8_t packetEncryptionReady(void) {
	return _encryptionReady;
}


void packetEncryptionSetup(PacketThreadDataT *threadData) {
	PacketEncodeDataT encoded                             = { 0 };
	uint16_t          dhData[PACKET_ENCRYPT_KEY_SIZE * 3] = { 0 };
	uint8_t           x                                   = 0;

	// Only call this from ONE SIDE of the connection.  It sets up both sides.

	for (x=0; x<PACKET_ENCRYPT_KEY_SIZE; x++) {
		threadData->dhModulus[x]  = PRIMES[rand() % PRIME_COUNT];
		threadData->dhBase[x]     = (rand() < (RAND_MAX / 2) ? 2 : 5);
		threadData->dhMySecret[x] = rand();
		threadData->dhMyPublic[x] = packetDHCompute(threadData->dhBase[x], threadData->dhMySecret[x], threadData->dhModulus[x]);
	}

	memcpy(&dhData[0],                           threadData->dhModulus,  PACKET_ENCRYPT_KEY_SIZE * sizeof(uint16_t));
	memcpy(&dhData[PACKET_ENCRYPT_KEY_SIZE],     threadData->dhBase,     PACKET_ENCRYPT_KEY_SIZE * sizeof(uint16_t));
	memcpy(&dhData[PACKET_ENCRYPT_KEY_SIZE * 2], threadData->dhMyPublic, PACKET_ENCRYPT_KEY_SIZE * sizeof(uint16_t));

	encoded.control    = PACKET_CONTROL_DAT;
	encoded.packetType = PACKET_TYPE_DH_REQUEST;
	encoded.channel    = 0;  // Doesn't matter for DH_REQUEST.
	encoded.encrypt    = 0;  // Must be 0 for DH_REQUEST.
	packetEncode(threadData, &encoded, (char *)dhData, PACKET_ENCRYPT_KEY_SIZE * 3 * sizeof(uint16_t));
	packetSend(threadData, &encoded);
}


void packetSend(PacketThreadDataT *threadData, PacketEncodeDataT *data) {
	// Valid control type?
	if (data->control != PACKET_CONTROL_BAD && data->control <= PACKET_CONTROL_COUNT) {
		_packetSender(data->dataPointer, data->length, threadData->senderData);
	} else {
		logWrite("Invalid PACKET_CONTROL!\n\r");
	}

	/*
	// Add to history?
	if (data->control == PACKET_CONTROL_DAT) {
		// ***TODO*** Must change control to use CONTROL_RTX & fix framing changes
		threadData->history[threadData->historyPosition].sequence = data->sequence;
		threadData->history[threadData->historyPosition].length   = data->length;
		memcpy(threadData->history[threadData->historyPosition].data, data->dataPointer, data->length);
		threadData->historyPosition++;
		if (threadData->historyPosition >= PACKET_SEQUENCE_MAX) {
			threadData->historyPosition = 0;
		}
	}
	*/

	// Mark invalid so caller has to change it.
	data->control = PACKET_CONTROL_BAD;
}


void packetSenderRegister(packetSender sender) {
	_packetSender = sender;
}


PacketThreadDataT *packetThreadDataCreate(void *senderData) {
	PacketThreadDataT *data = NULL;

	data = (PacketThreadDataT *)malloc(sizeof(PacketThreadDataT));
	if (data) {
		/*
		data->sequence           = 0;
		data->lastRemoteSequence = 0;
		data->historyPosition    = 0;
		*/
		data->decodeQueueHead    = 0;
		data->decodeQueueTail    = 0;
		data->newPacket          = 1;
		data->senderData         = senderData;
		data->aesContext         = NULL;
	}

	return data;
}


void packetThreadDataDestroy(PacketThreadDataT **data) {
	PacketThreadDataT *d = *data;

	if (d->aesContext) free(d->aesContext);
	free(d);
	d = NULL;
	*data = d;
}