/*

Regarding lincense information of any cipher algorithm used in the tests
please refer to:

RC4  		Ron Rivest (RSA Security)

MD5			Ron Rivest (RSA Security)

RC4PP		Shevek at https://crypto.anarres.info/2016/rc4_pencilandpaper

zx8			Karl-Uwe Frank at http://www.freecx.co.uk/crypto/zx8/

SBox8		Karl-Uwe Frank at http://www.freecx.co.uk/crypto/SBox8/

----------------------------------------------------------------------------


The test routine made by Karl-Uwe Frank and licensed as:

This is free and unencumbered software released into the public domain.

Anyone is free to copy, modify, publish, use, compile, sell, or
distribute this software, either in source code form or as a compiled
binary, for any purpose, commercial or non-commercial, and by any
means.

In jurisdictions that recognize copyright laws, the author or authors
of this software dedicate any and all copyright interest in the
software to the public domain. We make this dedication for the benefit
of the public at large and to the detriment of our heirs and
successors. We intend this dedication to be an overt act of
relinquishment in perpetuity of all present and future rights to this
software under copyright law.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
IN NO EVENT SHALL THE AUTHORS BE LIABLE FOR ANY CLAIM, DAMAGES OR
OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
OTHER DEALINGS IN THE SOFTWARE.

For more information, please refer to <http://unlicense.org/>

*/

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
// needed to generate uint32 on a 64bit OS
#include <stdint.h>

#include "md5/md5.h"
#include "md5/md5.c"


#define swap(X,Y) { uint8_t T = X; X = Y; Y = T; }

// rm -f ./keystream_byte_positions; gcc -std=c99 -O3 -fomit-frame-pointer -mtune=native keystream_byte_positions.c -o keystream_byte_positions

// for i in {aaa,bbb,ccc,ddd}; do nice -n 19 ./keystream_byte_positions 128 256 $((8**8)) $i rc4; done

// for i in {64,128,256}; do nice -n 19 ./keystream_byte_positions $i 256 $((8**8)) abc rc4; done


void MD5hash(unsigned char *data, unsigned int dataLen, unsigned char *digest) {
    MD5_CTX c;
    MD5Init(&c);
    MD5Update(&c, data, dataLen);
    MD5Final(digest, &c);
}


static uint8_t a, b, c, d, j, k, t;
static uint8_t SBox[256];
static uint8_t z[256], x[256];

void rc4_KSA(uint8_t key[], uint8_t keyLen);
void rc4x_KSA(uint8_t key[], uint8_t keyLen);
uint8_t rc4_PRGA();
uint8_t rc4_PP();

void SBox8_INIT();
void SBox8_KSA(uint8_t key[], uint8_t keyLen);
uint8_t SBox8_PRGA();
uint8_t SBox8a2_PRGA();

void zx8_KSA(uint8_t key[], uint8_t keyLen);
uint8_t zx8_PRGA();


//#############################################################


//-----------------------------------------
// KSA = Key Schedule Algorithm
//
void rc4_KSA(uint8_t key[], uint8_t keyLen) {
  uint8_t j = k = 0;

  for (int i=0; i<256; i++) {
    SBox[i] = i;
  }
    
  for (int i=0; i<256; i++) {
    k = (i % keyLen);
    j = (j + SBox[i] + key[k]);
    swap(SBox[i], SBox[j]);
  }
    
  a=0; b=0;
}


//-----------------------------------------
// KSA = Key Schedule Algorithm
//
void rc4x_KSA(uint8_t key[], uint8_t keyLen) {
  uint8_t j = k = 0;

  for (int i=0; i<256; i++) {
    SBox[i] = i;
  }
    
  for (int i=0; i<256; i++) {
    k = (i % keyLen);

    for (int n=0; n<8; n++) j = rc4_PRGA();
    j = j + SBox[i] + key[k];
    swap(SBox[i], SBox[j]);
  }
    
  a=0; b=0;
}


//-----------------------------------------
//
// Pseudo Random Generation Algorithm
//
uint8_t rc4_PRGA() {
  uint8_t t = 0;

  a = a + 1;
  b = b + SBox[a];
  
  swap(SBox[a], SBox[b]);
  
  t = SBox[a] + SBox[b];
  
  return SBox[t];
}



//#############################################################


//-----------------------------------------
//
// Pen and Paper RC4 as proposed at 
// https://crypto.anarres.info/2016/rc4_pencilandpaper
//
uint8_t rc4_PP() {
  uint8_t t = 0;
  uint8_t k = b;

  a = a + 1;
  b = b + SBox[a];
  
  swap(SBox[a], SBox[b]);
  
  t = SBox[a] + SBox[b] + k;
  
  return SBox[t];
}


//#############################################################

//-----------------------------------------
// INIT = Initialise the Internal State
//
void SBox8_INIT() {
	// SBox initalisation
  for (int i=0; i<256; i++) {
  	SBox[i] = i;
	}

	a = b = c = d = j = 0;
}


//-----------------------------------------
// KSA = Key Schedule Algorithm (Version 2)
//
void SBox8_KSA(uint8_t key[], uint8_t keyLen) {
	uint8_t t, k = 0;
	
	// Shuffle the Key into the Permutation
	for (int i=0; i<256; i++) {
		k  = i % keyLen;
		
    for (int n=0; n<8; n++) t = SBox8_PRGA();
    t = t + SBox[i] + key[k];
    swap(SBox[i], SBox[t]);
	}
}



//-----------------------------------------
// KSA = Key Schedule Algorithm
//
void SBox8a2_KSA(uint8_t key[], uint8_t keyLen) {
	uint8_t k = 0;
	
	// Shuffle the Key into the Permutation
	for (int i=0; i<256; i++) {
		k  = i % keyLen;
		
    for (int n=0; n<8; n++) j = SBox8_PRGA();
    j = j + SBox[i] + key[k];
    swap(SBox[i], SBox[j]);
	}
}


//-----------------------------------------
//
// Pseudo Random Generation Algorithm (Version 2)
//
uint8_t SBox8_PRGA() {
	a += SBox[b];
	b += SBox[a];
	c += (a ^ b) + j + SBox[j];
	
	swap(SBox[c], SBox[j]);		

	d  = SBox[(SBox[c] + SBox[d]) &0xff];
	
	j++;

	return ((a + b) ^ (c + d));
}


//-----------------------------------------
//
// Pseudo Random Generation Algorithm (Version a2)
//
uint8_t SBox8a2_PRGA() {

	a = a + (SBox[d] >> 5);
	b = a + (b ^ c);
	c = a + (SBox[b] << 3);
	d = a + (d ^ t);
	t = a + t;

	swap(SBox[t], SBox[j]);		

	j++;

	return (b ^ c ^ d);
}


//#############################################################


//-----------------------------------------
// 
// Key Schedule Algorithm (ps)
//
void zx8_KSA(uint8_t key[], uint8_t keyLen) {
  // Key must be at least 12 Characters long
  // and should have >= 60-Bit of Entropy.
  int i;
  unsigned char j, k, n, t; 

  // Prefill the Arrays
  for (i=0; i<256; i++) {
    z[i] = i;
    x[i] = i;
  }
   
  a=0; b=0; j=0; k=0; n=0, t=0;
  
  for (i=0; i<256; i++) {
    k = (i % keyLen);

    for (n=0; n<128; n++) t = zx8_PRGA();
    j = (t + j + z[i] + key[k]);
    swap(z[i], z[j]);
    
    for (n=0; n<128; n++) t = zx8_PRGA();
    j = (t + j + x[i] + z[x[j]]);
    swap(x[i], x[j]);
  } 

  // Reset the Array Indices Start Point
  a=0; b=0;
}


//-----------------------------------------
//
// Pseudo Random Generation Algorithm (ps)
//
unsigned char zx8_PRGA() {
  unsigned char n1, n2, y, m;

  // Calculate distant Array Element Indices
  n1 = z[a] + x[a];
  n2 = z[b] + x[b];

  // First Swap randomly selected Array Element
  swap(z[a], z[n1]);
  swap(x[a], x[n2]);

  // Update the global Carry on Array Indices
  a = a + b + (n1^n2);
  b = b + 1;

  // Second Swap sequentially cycle over every Array Element
  swap(z[b], z[n1]);
  swap(x[b], x[n2]);

  // Calculate the internal State Selector Value
  y = (z[n1] ^ x[n2]);

  // Calculate the internal State Protection Value
  m = (n1 + n2);

  // Never reveal internal State Values directly
  return (z[x[y]] ^ m);
}



//#############################################################

//-----------------------------------------
//
void ShowUsage(char* ThisName) 
{
	fprintf(stderr, "\nUsage  : %s 128 256 65536 abc rc4 (0.07)\n", ThisName);
	fprintf(stderr, "\n         128 = Size of Keys in Bit (32 to 256)\n");
	fprintf(stderr, "\n         256 = Length of Keystream in Byte (1 to 256)\n");
	fprintf(stderr, "\n         abc = Initial Value for Hash Generation\n");
	fprintf(stderr, "\n       65536 = Quantity of Keystreams to generate\n");
	fprintf(stderr, "\n         rc4 = Algorithm for Keystream Generation (rc4, rc4pp, sbox8, zx8)\n");
	fprintf(stderr, "\n");
	fprintf(stderr, "\nOptional:");
	fprintf(stderr, "\n         0.07 = Percentage of upper and lower Limit per Occurence of a Byte\n");
}


int main (int argc, char *argv[]) {

	uint16_t Key_BitSize 		= 128;
	uint16_t Key_ByteSize 	= Key_BitSize/8;
	uint32_t KeyStream_Len 	= 256;
	
	uint64_t TestRunCount 	= 65536;
	uint64_t OverflowLimit	= 0;
	uint64_t UnderflowLimit	= 0;
	uint64_t OverflowMax		= 0;
	uint64_t UnderflowMin		= 0;
	uint64_t Overflow				= 0;
	uint64_t Underflow			= 0;
	
	float  	LimitPercent		= 0.10;

	uint8_t  out_8bit				= 0;
		
	uint8_t CipherSwitch	  = 0;
	char HashInit[64] 			= "\0";
	char CipherName[10]			= "\0";

	char hashString[128] 		= "\0";
	char tempString[128] 		= "\0";

	uint8_t hashDigest[256] = {0,};
	uint8_t tempDigest[256] = {0,};

  uint32_t byte_matrix[256][256];
  uint32_t byte_homebox[256];

  memset(byte_matrix,  0, sizeof(uint32_t)*256*256);
  memset(byte_homebox, 0, sizeof(uint32_t)*256);
	memset(HashInit, 	 '\0', strlen(HashInit));
	memset(CipherName, '\0', strlen(CipherName));
	memset(hashString, '\0', strlen(hashString));
	memset(tempString, '\0', strlen(tempString));

  char* PrgName;

  PrgName = argv[0];

 	if (argc < 2) {
 		ShowUsage(PrgName); 
 		return 1;
 	}

 	if (argc >= 5) {
		Key_BitSize 		= atoi(argv[1]);
		Key_ByteSize 		= Key_BitSize/8;
		KeyStream_Len 	= atoi(argv[2]);
		TestRunCount 		= atoi(argv[3]);
	 	snprintf(HashInit, sizeof(HashInit), "%s", argv[4]);

 		if ((Key_BitSize == 0) || (Key_BitSize > 256) || (KeyStream_Len == 0) || (KeyStream_Len > 256) || (TestRunCount == 0)) {
 			ShowUsage(PrgName); 
 			return 1;
 		}
 	}	

 	
	// If a Limit Value is passed keep it
	if ((argc == 7) && (atof(argv[6]) > 0)) {
		LimitPercent = atof(argv[6]);
	} else {
		// Adjust the Base Limit Value
		uint64_t temp		= (TestRunCount / 256);
		if (temp > 256) {
			LimitPercent 		= (LimitPercent/2);
			OverflowLimit		= temp * (1 + LimitPercent);
			UnderflowLimit	= temp / (1 + LimitPercent);
		} else {
			// Lift the Limit
			LimitPercent		= 0.20;
			OverflowLimit		= 256 * (1 + LimitPercent);
			UnderflowLimit	= 256 / (1 + LimitPercent);
		}
	}
	
	OverflowMax			= OverflowLimit;
	UnderflowMin		= UnderflowLimit;
 	
 	if (argc >= 6) {
	 	snprintf(CipherName, sizeof(CipherName), "%s", argv[5]);
	}

	if (strcmp(CipherName, "sbox8") == 0) {
		CipherSwitch = 1;
		}
	else if (strcmp(CipherName, "zx8") == 0) {
		CipherSwitch = 2;
		}	
	else if (strcmp(CipherName, "rc4pp") == 0) {
		CipherSwitch = 3;
		}	
	else if (strcmp(CipherName, "rc4x") == 0) {
		CipherSwitch = 4;
		}	
	else if (strcmp(CipherName, "sbox8a2") == 0) {
		CipherSwitch = 5;
		}	
	else { 
		CipherSwitch = 0;
	 	snprintf(CipherName, sizeof(CipherName), "%s", "rc4\0");
	}


  char OutFile[129] = "";
  snprintf(OutFile, sizeof(OutFile), "bytepos_matrix_%s-%u-%u-%llu-%s.txt", CipherName, Key_BitSize, KeyStream_Len, TestRunCount, HashInit);
	OutFile[sizeof(OutFile)] = '\0';
	
  FILE* fOutMatrixTXT = fopen(OutFile, "w+");
  if (fOutMatrixTXT == NULL) { fprintf(stderr, "error creating file: %s\n", OutFile); return -1; }

  snprintf(OutFile, sizeof(OutFile), "bytepos_matrix_%s-%u-%u-%llu-%s.csv", CipherName, Key_BitSize, KeyStream_Len, TestRunCount, HashInit);
	OutFile[sizeof(OutFile)] = '\0';
	
  FILE* fOutMatrixCSV = fopen(OutFile, "w+");
  if (fOutMatrixCSV == NULL) { fprintf(stderr, "error creating file: %s\n", OutFile); return -1; }
	
  snprintf(OutFile, sizeof(OutFile), "byte_in_homebox_%s-%u-%u-%llu-%s.csv", CipherName, Key_BitSize, KeyStream_Len, TestRunCount, HashInit);
	OutFile[sizeof(OutFile)] = '\0';
	
  FILE* fOutHomeBoxCSV = fopen(OutFile, "w+");
  if (fOutHomeBoxCSV == NULL) { fprintf(stderr, "error creating file: %s\n", OutFile); return -1; }

	
	// Build the first Hash String for the Random Keys
	snprintf(hashString, sizeof(hashString), "%s", HashInit);
	uint32_t dLen = strlen(hashString);

	// Test Loop
	for (uint64_t runs=0; runs<TestRunCount; runs++) {
		uint8_t key[256];
		uint8_t keyLen = Key_ByteSize;
		memset(key, 0, 256);

		// printf("%llu) hashString = %s\n", runs, hashString);

		MD5hash(hashString, dLen, hashDigest);

		// Concatenate Hashes
		if (Key_BitSize > 128) {

			// Prepend Hex String if any
			for (int i=0; i<16; i++) {
				// Save the current Digest
				tempDigest[i] = hashDigest[i];
				// Build a String of the current Digest in Hex String Notation
				//				snprintf(tempString, sizeof(tempString), "%s%02x", tempString, hashDigest[i]);   // some compiler don't understand this
				sprintf(tempString, "%s%02x", tempString, hashDigest[i]);
			}
			tempString[sizeof(tempString)] = '\0';
			dLen = strlen(tempString);

			//--> printf("tempString = %s\n", tempString);

			MD5hash(tempString, dLen, hashDigest);

			// Extend the Digest
			for (int i=16; i<keyLen; i++) {
				hashDigest[i] = tempDigest[i-16];
			}			
		}
	
		memset(hashString, '\0', strlen(hashString));
		memset(tempString, '\0', strlen(tempString));
	
		// Pass Digest to Key
		for (int i=0; i<keyLen; i++) {
			key[i] = hashDigest[i];
			
			// Use the last Hash String in Hex String Notation for Generation of the next Key
			//			snprintf(hashString, sizeof(hashString), "%s%02x", hashString, key[i]);   // some compiler don't understand this
			sprintf(hashString, "%s%02x", hashString, key[i]);
		}
		hashString[sizeof(hashString)] = '\0';

		// Adjust the Hash String Length
		dLen = strlen(hashString);
			
		memset(hashDigest, 0, sizeof(hashDigest));
		memset(tempDigest, 0, sizeof(tempDigest));

		// Print the Key
  	//--> for (int i=0; i<keyLen; i++) printf("%02x", key[i]);  printf("\n");

		// Initialise the Stream Cipher
		switch (CipherSwitch) {
 			case 0: rc4_KSA(key, keyLen); break;
			case 1: SBox8_INIT(); SBox8_KSA(key, keyLen); break;
			case 2: zx8_KSA(key, keyLen); break;
			case 3: rc4_KSA(key, keyLen); break;
			case 4: rc4x_KSA(key, keyLen); break;
			case 5: SBox8_INIT(); SBox8a2_KSA(key, keyLen); break;
		}

		// Fill the Keystream into the Matrix
		for (int i=0; i<KeyStream_Len; i++) {
			out_8bit = 0;
			
			switch (CipherSwitch) {
				case 0: out_8bit = rc4_PRGA(); break;
				case 1: out_8bit = SBox8_PRGA(); break;
				case 2: out_8bit = zx8_PRGA(); break;
				case 3: out_8bit = rc4_PP(); break;
				case 4: out_8bit = rc4_PRGA(); break;
				case 5: out_8bit = SBox8a2_PRGA(); break;
			}
		
			byte_matrix[i][out_8bit]++;
			byte_homebox[out_8bit]++;
		}
	}		 


	fprintf(fOutMatrixTXT, "Matrix of row Occurance in Keystream after the KSA has finished.\n");
	fprintf(fOutMatrixTXT, "=================================================================\n\n");
	fprintf(fOutMatrixTXT, "Key Bitsize            : %d Bit\n", Key_BitSize);
	fprintf(fOutMatrixTXT, "Keystream Length       : %d Byte\n", KeyStream_Len);
	fprintf(fOutMatrixTXT, "Random Keys Count      : %-16llu\n", TestRunCount);
	fprintf(fOutMatrixTXT, "Random Keys Startvalue : %s\n", HashInit);
	fprintf(fOutMatrixTXT, "Cipher Algorithm       : %s\n", CipherName);

	fprintf(fOutMatrixTXT, "\n\nCount of Byte at Position in Keystream \n\n");
	
	fprintf(fOutMatrixTXT, "     %10d ", 0);
	fprintf(fOutMatrixCSV, "\"Byte @ Pos\",0,");	    
	for (int row=1; row<KeyStream_Len; row++) {
		fprintf(fOutMatrixTXT, "%10d ", row);
		fprintf(fOutMatrixCSV, "%d,", row);	    
	}
	
	fprintf(fOutMatrixTXT, "\n   |===========|");	
	for (int row=1; row<KeyStream_Len; row++) {
		fprintf(fOutMatrixTXT, "==========|");
	}
	fprintf(fOutMatrixTXT, "\n");
	fprintf(fOutMatrixCSV, "\n");	    
	
	// Initialise
	OverflowMax 	= byte_matrix[0][0];
	UnderflowMin 	= byte_matrix[0][0];

	for (int row=0; row<256; row++) {
		if ( ((row % 10) == 0) && (row > 1) ) {
			for (int col=0; col<KeyStream_Len; col++) {
				if (col<1)	fprintf(fOutMatrixTXT, "   |.");
				fprintf(fOutMatrixTXT, "..........|");
			}
			fprintf(fOutMatrixTXT,  "\n");
			fprintf(fOutMatrixTXT,  "%02x | ", row);	
			fprintf(fOutMatrixCSV,  "\"0x%02x\",", row);	    
			fprintf(fOutHomeBoxCSV, "\"0x%02x\",", row);	    
			}
		else {
			fprintf(fOutMatrixTXT,  "%02x | ", row);	
			fprintf(fOutMatrixCSV,  "\"0x%02x\",", row);	    
			fprintf(fOutHomeBoxCSV, "\"0x%02x\",", row);	    
		}
		
		for (int col=0; col<KeyStream_Len; col++) {
			uint32_t byte_matrix_value = byte_matrix[col][row];
			
			if ( byte_matrix_value > OverflowLimit ){
				fprintf(fOutMatrixTXT, "*%9d|", byte_matrix_value);
				// fprintf(fOutMatrixTXT, "%10d|", byte_matrix_value);	    
				Overflow++;
				}
			else if ( byte_matrix_value < UnderflowLimit ){
				fprintf(fOutMatrixTXT, "-%9d|", byte_matrix_value);
				// fprintf(fOutMatrixTXT, "%10d|", byte_matrix_value);	    
				Underflow++;
			}
			else {
				fprintf(fOutMatrixTXT, "%10d|", byte_matrix_value);	    
			}
	
			if (byte_matrix_value > OverflowMax) 	OverflowMax 	= byte_matrix_value;
			if (byte_matrix_value < UnderflowMin) UnderflowMin 	= byte_matrix_value;

			fprintf(fOutMatrixCSV,  "%d,", byte_matrix_value);	    
		}
	
		fprintf(fOutMatrixTXT,  "\n");
		fprintf(fOutMatrixCSV,  "\n");	    
		fprintf(fOutHomeBoxCSV, "%d\n", byte_homebox[row]);	    
	}
	

	fprintf(fOutMatrixTXT, "\nTotal (Limit %d%%)\n============================\n", (int)(LimitPercent *100));
	fprintf(fOutMatrixTXT, "Upper Limit         = %llu\n", OverflowLimit);
	fprintf(fOutMatrixTXT, "Upper Limit reached = %llu times\n", Overflow);
	fprintf(fOutMatrixTXT, "Highest Value found = %llu\n\n", OverflowMax);

	fprintf(fOutMatrixTXT, "Lower Limit         = %llu\n", UnderflowLimit);
	fprintf(fOutMatrixTXT, "Lower Limit reached = %llu times\n", Underflow);
	fprintf(fOutMatrixTXT, "Lowest Value found  = %llu\n\n", UnderflowMin);


	fclose(fOutMatrixTXT);
	fclose(fOutMatrixCSV);
	fclose(fOutHomeBoxCSV);
	
	return 0;
}