/* TODO:
 * 	1.  Does a 1-bit change actually change the other bits?  More evaluation
 * 	2.  Add some kind of stream operating mode, like nonce and counter
 */

#include<stdio.h>
#include<math.h>
#include<stdint.h>
#include<unistd.h>
#include<stdlib.h>

void shift_data(uint8_t* data){
	uint8_t tmp = data[0];
	data[0] = data[1];
	data[1] = data[2];
	data[2] = data[3];
	data[3] = tmp;
}

void unshift_data(uint8_t* data){
	uint8_t tmp = data[3];
	data[3] = data[2];
	data[2] = data[1];
	data[1] = data[0];
	data[0] = tmp;
}

const int round_key_primes[] = {31, 37, 61, 83, 13, 17, 91, 59, 57, 7, 3, 101, 113, 23, 41, 43};

uint32_t make_round_key(uint32_t key, int round){
	uint32_t round_key_int;
	uint8_t *round_key = (uint8_t*)&round_key_int;

	round_key[0] = key * round_key_primes[round * 4 + 0]; 
	round_key[1] = key * round_key_primes[round * 4 + 1];
	round_key[2] = key * round_key_primes[round * 4 + 2];
	round_key[3] = key * round_key_primes[round * 4 + 3];
	fprintf(stderr, "Round key %d:  %x\n", round, round_key_int);
	return round_key_int;
}

void encrypt_block(void* data, uint32_t key){
	uint8_t *typed_data = (uint8_t*)data;
	uint32_t *typed_data_32 = (uint32_t*)data;

	// # permute data to try to mix in changes
	// # apply some sort of substitution
	// # apply the key (this could affect the other two operations)
	// # mix data from each byte into each other byte
	// Note:  We have to be able to undo anything we do!
	
	for(int round = 0; round < 4; round++){
		uint32_t round_key = make_round_key(key, round);
		fprintf(stderr, "%x\n", *typed_data_32);
		*typed_data_32 = *typed_data_32 ^ round_key;
		fprintf(stderr, "%x\n", *typed_data_32);
		typed_data[0] ^= typed_data[3];
		fprintf(stderr, "%x\n", *typed_data_32);
		typed_data[1] ^= typed_data[2];
		fprintf(stderr, "%x\n", *typed_data_32);
		shift_data(typed_data);
		fprintf(stderr, "%x\n", *typed_data_32);

	}
}

void decrypt_block(void* data, uint32_t key){
	uint8_t *typed_data = (uint8_t*)data;
	uint32_t *typed_data_32 = (uint32_t*)data;

	for(int round = 0; round < 4; round++){
		uint32_t round_key = make_round_key(key, 3 - round);
		unshift_data(typed_data);
		typed_data[0] ^= typed_data[3];
		typed_data[1] ^= typed_data[2];
		*typed_data_32 = *typed_data_32 ^ round_key;
	}
}

void encrypt_data(void* data, uint8_t key, size_t datasize, char decrypt){
	size_t offset = 0;
	while(offset < datasize){
		if(!decrypt)
			encrypt_block(data + offset, key);
		else
			decrypt_block(data + offset, key);
		offset += 4;
	}
}

#define BUFSIZE (1024*1024)
int main(int argc, char ** argv){
	uint32_t key = atol(argv[1]);

	uint8_t buffer[BUFSIZE];
	ssize_t readlen;
	for(;;){
		readlen = read(0, buffer, BUFSIZE);
		if(readlen < 1)
			break;
		encrypt_data(buffer, key, readlen, argc > 2);
		write(1, buffer, readlen);
	}
	
	return 0;
}