/****************************************************************/
/* 802.11 RC4 Test Vector Generator                             */
/* (c) 2002, David Johnston                                     */
/* Author: David Johnston                                       */
/* Email (home): dj@deadhat.com                                 */
/* Email (general): david.johnston@ieee.org                     */
/* Version 0.1                                                  */
/****************************************************************/
 
#include <stdlib.h>
#include <stdio.h>

#define MAX_MESSAGE_LENGTH 4096

/********************************************/
/* Test Cases                               */
/* An array of test cases, each containing  */
/* an MPDU. A separate array gives the      */
/* length of each MPDU.                     */
/********************************************/
#define NUM_TEST_CASES  3

  

int key_lengths[] = {8,16, 16};
unsigned char keys[] =
{
    0x5a,0xa5,0x12,0x34,0x56,0x99,0x88,0x77,
    0xc0,0xc1,0xc2,0xc3,0xc4,0xc5,0xc6,0xc7,0xc8,0xc9,0xca,0xcb,0xcc,0xcd,0xce,0xcf,
    0x5b,0x7b,0xa0,0xd7,0x9a,0xee,0xc2,0x2e,0x0d,0xd1,0xa9,0x14,0xbd,0xb8,0x42,0x30

};

int data_lengths[] = {41, 110, 51};
unsigned char datas[] =                                 /* Test Case 1 */
{
    0x08,0x01,0x02,0x01,0x00,0x06,0x25,0xa7,
    0xc4,0x36,0x00,0x02,0x2d,0x49,0x97,0xb4,
    0x00,0x06,0x25,0xa7,0xc4,0x36,0xe0,0x00,
    0xaa,0xaa,0x03,0x00,0x00,0x00,0x88,0x8e,
    0x01,0x01,0x00,0x00,0x00,0x00,0x00,0x00,
    0x00,
                                                         /* Test Case 2 */

    0x08,0x03,0x12,0x34,0xff,0xff,0xff,0xff,
    0xff,0xff,0x00,0x40,0x96,0x45,0x07,0xf1,
    0x08,0x00,0x46,0x17,0x62,0x3e,0x50,0x67,
    0xaa,0xaa,0x03,0x00,0x00,0x00,0x08,0x00,
    0x45,0x00,0x00,0x4e,0x66,0x1a,0x00,0x00,
    0x80,0x11,0xbe,0x64,0x0a,0x00,0x01,0x22,
    0x0a,0xff,0xff,0xff,0x00,0x89,0x00,0x89,
    0x00,0x3a,0x00,0x00,0x80,0xa6,0x01,0x10,
    0x00,0x01,0x00,0x00,0x00,0x00,0x00,0x00,
    0x20,0x45,0x43,0x45,0x4a,0x45,0x48,0x45,
    0x43,0x46,0x43,0x45,0x50,0x46,0x45,0x45,
    0x49,0x45,0x46,0x46,0x43,0x43,0x41,0x43,
    0x41,0x43,0x41,0x43,0x41,0x43,0x41,0x41,
    0x41,0x00,0x00,0x20,0x00,0x01,

    0xaa,0xaa,0x03,0x00,0x00,0x00,0x08,0x00,           /* Test Case 3 */
    0x45,0x00,0x00,0x4e,0x66,0x1a,0x00,0x00,
    0x80,0x11,0xbe,0x64,0x0a,0x00,0x01,0x22,
    0x0a,0xff,0xff,0xff,0x00,0x89,0x00,0x89,
    0x00,0x3a,0x00,0x00,0x80,0xa6,0x01,0x10,
    0x00,0x01,0x00,0x00,0x00,0x00,0x00,0x99,
    0x22,0x5f,0x4e
};





void rc4(
            unsigned char *key,
            int key_length,
            int cipherstream_length,
            unsigned char *cipherstream
        );

void rc4_encrypt(
            unsigned char *key,
            int key_length,
            unsigned char *data,
            int data_length,
            unsigned char *ciphertext
        );
/****************************************/
/* rc4_encrypt()                        */
/****************************************/
void xor_block(int length, unsigned char *a, unsigned char *b, unsigned char *out)
{
    int i;
    for (i=0;i<length; i++)
    {
        out[i] = a[i] ^ b[i];
    }
}

void swap(unsigned char *a, unsigned char *b)
{
    unsigned char tmp;

    tmp = *a;
    *a = *b;
    *b = tmp;
}

void rc4(
            unsigned char *key,
            int key_length,
            int cipherstream_length,
            unsigned char *cipherstream
        )
{
    int i,j,x;
    unsigned char s[256];
    unsigned char k[256];
    
    /* Create Key Stream */

    for (i=0;i<256;i++)
    {
        k[i] = key[i % key_length];
    }

    /* Initialize SBOX */
    for (j=0; j<256; j++) s[j] = j;

    /* Seed the SBOX */
    i = 0;
    for (j=0; j<256; j++) 
    {
        i = (i + s[j] + k[j]) % 256;
        swap(&s[j],&s[i]);
    }

    /* Generate the cipherstream */
    j = 0;
    i = 0;

    for (x=0;x<cipherstream_length;x++)
    {
        j = (j +1)  % 256;
        i = (i + s[j])  % 256;
        swap(&s[j],&s[i]);
        cipherstream[x] = s[(s[j] + s[i])%256];
    };

}

void rc4_encrypt(
            unsigned char *key,
            int key_length,
            unsigned char *data,
            int data_length,
            unsigned char *ciphertext
        )

{
    unsigned char cipherstream[MAX_MESSAGE_LENGTH+1];
    rc4(key, key_length, data_length, cipherstream);

    xor_block(data_length, cipherstream, data, ciphertext);
}
/****************************************/
/* debug_out()                          */
/****************************************/

void debug_out(
            unsigned char *label,
            unsigned char *data,
            int data_length
        )
{
int i,j;
int num_blocks;
int block_remainder;
    num_blocks = data_length / 16;
    block_remainder = data_length % 16;

    printf("%s\n",label);

    for (i=0; i< num_blocks;i++)
    {
        printf("\t");
        for (j=0; j< 16;j++)
        {
            printf("%02x ", data[j + (i*16)]);
        }
        printf("\n");
    }

    if (block_remainder > 0)
    {
        printf("\t");
        for (j=0; j<block_remainder; j++)
        {
            printf("%02x ", data[j+(num_blocks*16)]);
        }
        printf("\n");
    }
}

/************************************************/
/* get_testcase()                               */
/* Copies data and keys from test cases         */
/************************************************/
void get_testcase(
                    int test_case,
                    unsigned char *plaintext,
                    int *length,
                    unsigned char *key,
                    int *key_length
                )
{

    int i;
    unsigned char *ptr;

    ptr = datas;
    for (i=0; i< (test_case-1); i++)    /* Iterate through test cases */
    {
        ptr = ptr + data_lengths[i];
    }

    for (i=0; i< data_lengths[test_case-1]; i++)
    {
        plaintext[i] = *ptr++;
    }
    *length = data_lengths[test_case-1];


    ptr = keys;
    for (i=0; i< (test_case-1); i++)    /* Iterate through test cases */
    {
        ptr = ptr + key_lengths[i];
    }

    for (i=0; i< key_lengths[test_case-1]; i++)
    {
        key[i] = *ptr++;
    }
    *key_length = key_lengths[test_case-1];
}
/****************************************************/
/* main()                                           */
/* Iterate through the test cases, passing them     */
/* through the ccm algorithm to produce test        */
/* vectors                                          */
/****************************************************/

int main()
{
    int length;
    int key_length;
    int test_case;

    unsigned char plaintext[MAX_MESSAGE_LENGTH+1];
    unsigned char cipherstream[MAX_MESSAGE_LENGTH+1];
    unsigned char ciphertext[MAX_MESSAGE_LENGTH+1];
    unsigned char key[256];

    for (test_case = 1; test_case < (NUM_TEST_CASES+1); test_case++)
    {
        printf ("\nTEST CASE %d\n",test_case);

        get_testcase(test_case, plaintext, &length, key, &key_length);

        debug_out("Key       = ",key, key_length);
        debug_out("Plaintext = ",plaintext, length);

        rc4(key,key_length,length,cipherstream);
        
        debug_out("Cipherstream = ",cipherstream, length);

        rc4_encrypt(key,key_length,plaintext,length,ciphertext);

        debug_out("Ciphertext = ",ciphertext, length);
    }

    return 0;
}