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


/* It has been modified to compute a histogram of byte values in the SBox[]
 * array after the KSA has been run */

// gcc -std=c99 -O3 -fomit-frame-pointer -mtune=native rc4.testksa2.c -o rc4.testksa2


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

static uint8_t  a, b, k, j;
static uint8_t  SBox[256];

void RC4_KSA(uint8_t key[], uint8_t keyLen) {
	for (int i=0; i<256; i++) {
    SBox[i] = i;
  }

	k=0; j=0;
	
	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;
} 

/* Test harness code below */

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

  /* 75, 124, 88 were picked at random */
  uint8_t key[4] = {75, 0, 0, 0};

  unsigned int sbox_histogram[256][256];
  memset(&sbox_histogram, 0, sizeof(int)*256*256);

	uint16_t maxdeeps = 256;

  for (int y=0; y<256; y++) {
    key[1] = y & 0xff;
    for (int x=0; x<256; x++) {
      key[2] = x & 0xff;
      for (int i=0; i<256; i++) {
        key[3] = i & 0xff;
        RC4_KSA(key, 4);
        for (int idx=0; idx<256; idx++) {
          sbox_histogram[idx][SBox[idx]]++;
        }
      }
    }  
	}

  printf("Histogram of SBox[] after KSA for key[2-3] from 0x0 to 0xffff:\n");
  printf("       SBox[i]:\n     ");
  for (int i=0; i<256; i++) {
    printf("%10d ", i);
  }
  printf("\nVal: ");
   for (int i=0; i<256; i++) {
    printf("===========");
   }
   printf("\n");
   for (int byte=0; byte<256; byte++) {
     printf("%02x:  ", byte);
     for (int idx=0; idx<256; idx++) {
     	if ( sbox_histogram[idx][byte] > 100000 ){
	      printf("***%7d ", sbox_histogram[idx][byte]);
	    	}
     	else if ( sbox_histogram[idx][byte] < 24000 ){
	      printf("---%7d ", sbox_histogram[idx][byte]);
	    	}
	    else {
	      printf("%10d ", sbox_histogram[idx][byte]);	    
	    }
     }
     printf("\n");
   }
}