/* RC4 Book Stack Test Output implemented by Karl-Uwe Frank */

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

typedef struct {
  unsigned char SBox[256];
  unsigned char i, j;
} RC4_CTX;


/* Function prototypes */
void RC4_set_key(RC4_CTX *c, size_t keybytes, unsigned char *key);
void RC4(RC4_CTX *c, size_t n, unsigned char *in, unsigned char *out);


//-----------------------------------------
// 
// rm -f rc4__BookStack && gcc -O3 -std=c99 rc4__BookStack.c -o rc4__BookStack
// 
// ./rc4__BookStack 11 256 65432912 > rc4_BookStack.out
// 


//-----------------------------------------
//
void ShowUsage(char* ThisName) 
{
	fprintf(stderr, "\nProgram usage:\n");
	fprintf(stderr, "%s 11  256  65432912\n\n", ThisName);
	fprintf(stderr, "       where in the example above means\n\n");
	fprintf(stderr, "      11 = 2**11 test bit length (minimum 2**11)\n");
	fprintf(stderr, "     256 = key bit size\n");
	fprintf(stderr, "65432912 = Seed (for PRNG to genereate test keys)\n");
	fprintf(stderr, "           if Seed = 0 then timestamp will be used\n\n");
}

//-----------------------------------------
// Main 
//
int main (int argc, char *argv[]) 
{ 
  // initialisation
  char* PrgName;
	unsigned int Seed; 

	time_t Seed_Time;
	Seed_Time = time (NULL);
	Seed			= Seed_Time;

	uint64_t TestSize = (pow(2,24)/8); // bit to byte
	int KeyBitSize = 256;

  const    int  BUF_SIZE = 16384;
  unsigned char inB[16384] = {0,};
	unsigned char outB[16384];


	PrgName = argv[0];
	
	if (argc > 1) {
		if ( (strcmp(argv[1], "-h") == 0) || (strcmp(argv[1], "--help") == 0) ) {
			ShowUsage(PrgName);
			exit(1);
		}
	
		if ( atoi(argv[1]) < 11) {
			ShowUsage(PrgName);
			exit(1);
		}
	
		if ( argc > 1) TestSize = (pow(2, atoi(argv[1]))/8);
		if ( argc > 2) KeyBitSize = atoi(argv[2]);
		if ((argc > 3) && (atoi(argv[3]) > 0)) Seed = atoi(argv[3]);
	}		

  //-----------------------------------------
  //
  // Set the Encryption Keyword
  //
	int i, KeyLen;

	unsigned char KeyWord[256] = {0,};

//---------------------------------------

	srand( Seed );
	for (i=0; i<(KeyBitSize/8); i++) KeyWord[i] = rand() % 256;
	KeyWord[i] = '\0';
	KeyLen 		 = i-1;
	


  //-----------------------------------------
  // Initialisation of rc4
  //
  // perform the Key Schedule
	RC4_CTX ctx;
	RC4_set_key(&ctx, KeyLen, KeyWord);


// 	fprintf (stderr,"=========================================================================================================\n");
// 	fprintf (stderr,"Book Stack Test : RC4\n");
// 	fprintf (stderr,"=========================================================================================================\n");
//   fprintf (stderr,"Keyword is      : "); for (i=0; i<KeyLen; i++) fprintf (stderr,"%02x", KeyWord[i]); fprintf (stderr,"\n");
//   fprintf (stderr,"MB Test Size    : %10llu\n", (TestSize/1024/1024));
// 	fprintf (stderr,"=========================================================================================================\n");


  // Write to STDOUT
  for (uint64_t i=0; i < TestSize; i++) {
    RC4(&ctx, 1, inB, outB);
    fwrite(&outB, sizeof(char), 1, stdout);    
	}

  return 0;
} 



//
// RC4 Source Code found at
//
// http://etutorials.org/Programming/secure+programming/Chapter+5.+Symmetric+Encryption/5.23+Setting+Up+and+Using+RC4/
//

void RC4_set_key(RC4_CTX *c, size_t keybytes, unsigned char *key) {
  int           i, j;
  unsigned char keyarr[256], swap;
   
  c->i = c->j = 0;
  for (i = j = 0;  i < 256;  i++, j = (j + 1) % keybytes) {
    c->SBox[i] = i;
    keyarr[i] = key[j];
  }
  for (i = j = 0;  i < 256;  i++) {
    j += c->SBox[i] + keyarr[i];
    j %= 256;
    swap = c->SBox[i];
    c->SBox[i] = c->SBox[j];
    c->SBox[j] = swap;
  }
}


void RC4(RC4_CTX *c, size_t n, unsigned char *in, unsigned char *out) {
  unsigned char swap;
   
  while (n--) {
    c->j += c->SBox[++c->i];
    swap = c->SBox[c->i];
    c->SBox[c->i] = c->SBox[c->j];
    c->SBox[c->j] = swap;
    swap = c->SBox[c->i] + c->SBox[c->j];
    *out++ = *in++ ^ c->SBox[swap];
  }
}