/* 
 * The following is an implementation of the ChaCha20 symmetric stream cipher algorithm for up to 256 ASCII character input in C.
 * By: Cruz Alcorta
 * Written with the help of this article plus Wikipedia and Google: https://loup-vaillant.fr/tutorials/chacha20-design
 */

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>

/* Performs a left bit rotation by "n" which loops the bits around to the right (rather than just dropping them off a cliff and leaving 0s. */
#define ROT_L32(x, n) (x = (x << n) | (x >> (32 - n)))

/* Performs a quarter-round to mix up four 32-bit integers */
#define QR(a, b, c, d)                    \
	do {                                  \
		a += b;  d ^= a;  ROT_L32(d, 16); \
		c += d;  b ^= c;  ROT_L32(b, 12); \
		a += b;  d ^= a;  ROT_L32(d,  8); \
		c += d;  b ^= c;  ROT_L32(b,  7); \
	} while (0)

void
complete_round(uint32_t *blk);

void
psuedo_random_block(uint32_t *blk);

int
main(int argc, char **argv)
{
	/* This is just the string "expand 32-byte k" represented as 4 chunks of uint32_t integers */
	uint32_t constant[4] = {
		0x65787061, 0x6E642033, 0x322D6279, 0x7465206B
	};

	uint64_t counter = 0;

	/* Get key and nonce from the user */
	uint32_t key[8];
	uint32_t nonce[2];

	/* This is only constant for now */
	uint32_t text[64]; // <---- ONLY CONSTANT FOR NOW

	/* The extra character is for the null-terminator. */
	char hexKey[65];
	char hexNonce[17];

	/* Max length of 256 character messages (1 ASCII character corresponds to 2 hex digits) */
	char hexText[513]; // <---- ONLY CONSTANT FOR NOW

	printf("KEY (256-bit hex string, 0-9 and A-F, 64 characters): ");
	scanf("%64s", hexKey);
	printf("NONCE (64-bit hex string, 0-9 and A-F, 16 characters): ");
	scanf("%16s", hexNonce);
	printf("INPUT (hex string, maximum of 512 characters will be read -> 256 ASCII characters): ");
	scanf("%512s", hexText);

	/* Basic sanity check on the user input */
	if (strlen(hexKey) != 64 || strlen(hexNonce) != 16) {
		printf("ERROR: size of hexadecimal key and nonce is too short.\n");
		return 1;
	}
	if (strspn(hexKey, "0123456789abcdefABCDEF") != 64 || strspn(hexNonce, "0123456789abcdefABCDEF") != 16 || strspn(hexText, "0123456789abcdefABCDEF") != strlen(hexText)) {
		printf("ERROR: key, nonce, or input are not in a valid hexadecimal format.\n");
		return 1;
	}

	/* Pad the hexText with '0' to match its full size */
	/* i < 512, as we don't want to overwrite the last character (the null terminator) */
	for (size_t i = strlen(hexText); i < 512; i++) { // <---- ONLY CONSTANT FOR NOW
		hexText[i] = '0';
	}

	/* Copy groups of 8 characters (each character representing one hexadecimal digit, or 4 bits, thus there are 32 bits in total) into the key, nonce, and text uint32_t vectors */
	for (int i = 0; i < 8; i++) {
		/* The extra character is for the null-terminator. */
		char tmp[9] = {0};
		strncpy(tmp, &hexKey[i * 8], 8);
		sscanf(tmp, "%x", &key[i]);
	}
	for (int i = 0; i < 2; i++) {
		/* The extra character is for the null-terminator. */
		char tmp[9] = {0};
		strncpy(tmp, &hexNonce[i * 8], 8);
		sscanf(tmp, "%x", &nonce[i]);
	}
	for (int i = 0; i < 64; i++) {
		/* The extra character is for the null-terminator. */
		char tmp[9] = {0};
		strncpy(tmp, &hexText[i * 8], 8);
		sscanf(tmp, "%x", &text[i]);
	}

	/* AT THIS POINT: the constant, key, nonce, and counter have all been initialized.  Now, generate the blocks and scramble them. */

	/* Generate 16 blocks */
	size_t NUM_BLOCKS = 16; // <---- ONLY CONSTANT FOR NOW

	/* The 2-dimensional array "blocks" holds "n" pointers to the start of seperate 512-bit blocks of 16 uint32_t chunks */
	uint32_t **blocks = calloc(NUM_BLOCKS, sizeof(uint32_t *));

	for (size_t i = 0; i < NUM_BLOCKS; i++) {
		blocks[i] = (uint32_t *)calloc(16, sizeof(uint32_t));

		/* Initialize the block */
		blocks[i][0] = constant[0];
		blocks[i][1] = constant[1];
		blocks[i][2] = constant[2];
		blocks[i][3] = constant[3];

		blocks[i][4]  = key[0];
		blocks[i][5]  = key[1];
		blocks[i][6]  = key[2];
		blocks[i][7]  = key[3];
		blocks[i][8]  = key[4];
		blocks[i][9]  = key[5];
		blocks[i][10] = key[6];
		blocks[i][11] = key[7];

		/* Split the lower and upper parts of the counter respectively into two 32-bit words*/
		blocks[i][12] = (uint32_t)(counter);
		blocks[i][13] = (uint32_t)(counter >> 32);

		blocks[i][14] = nonce[0];
		blocks[i][15] = nonce[1];

		/* Perform scrambling to turn it into a psuedo-random block of data */
		psuedo_random_block(blocks[i]);

		/* Increment the counter for the next block */
		counter++;
	}

	/* XOR our psuedo-random chunks (inside of the blocks) with the input text */
	for (size_t i = 0; i < 64; i++) {
		text[i] ^= blocks[i / NUM_BLOCKS][i % 16];
	}

	/* Convert it back to hexadecimal text (just reusing the hexText variable here) */
	char *textPos = hexText;
	for (size_t i = 0; i < 64; i++) {
		textPos += sprintf(textPos, "%08x", text[i]);
	}

	printf("RESULT: %s\n", hexText);

	/* Clean up the allocated memory */
	for (size_t i = 0; i < NUM_BLOCKS; i++) {
		free(blocks[i]);
	}
	free(blocks);

	return 0;
}

void
complete_round(uint32_t *blk)
{
	/* Note: we are assuming the 4x4 matrix that we are representing with a 16-dimensional vector has it's elements ordered like so, based on the 0-based index in the vector:
	 * 0   1   2   3
	 * 4   5   6   7
	 * 8   9   10  11
	 * 12  13  14  15
	 */
	for (int i = 0; i < 10; i++) {
		/* Perform a quarter-rotation on each of the 4 columns of the block */
		QR(blk[0], blk[4], blk[8],  blk[12]);
		QR(blk[1], blk[5], blk[9],  blk[13]);
		QR(blk[2], blk[6], blk[10], blk[14]);
		QR(blk[3], blk[7], blk[11], blk[15]);

		/* Perform a quarter-rotation on the 4 "diagonals" (each diagonal that runs out of room loops around to the opposite side of the matrix */
		QR(blk[0], blk[5], blk[10], blk[15]);
		QR(blk[1], blk[6], blk[11], blk[12]);
		QR(blk[2], blk[7], blk[8],  blk[13]);
		QR(blk[3], blk[4], blk[9],  blk[14]);
	}
}

void
psuedo_random_block(uint32_t *blk)
{
	uint32_t tmp[16];
	memcpy(tmp, blk, sizeof(tmp));

	complete_round(blk);

	for (size_t i = 0; i < 16; i++) {
		blk[i] += tmp[i];
	}
}