/****************************************************************/
/* 802.11 TKIP Michael Test Vector Generator                    */
/* (c) 2002, David Johnston                                     */
/* Author: David Johnston                                       */
/* Email (home): dj@deadhat.com                                 */
/* Email (general) : david.johnston@ieee.org                    */
/* Version 0.2                                                  */
/* v0.2                                                         */
/*     Fixed bug for longer messages                            */
/* v0.3                                                         */
/*     Fixed alignment bug in format_,essage()                  */
/****************************************************************/
 
#include <stdlib.h>
#include <stdio.h>

/****************************************/
/* Text output options.                 */
/* Comment out to reduce text output    */
/****************************************/


/********************************************/
/* Test Cases                               */
/* An array of test cases, each containing  */
/* a block to be processed by michael       */
/********************************************/
#define NUM_TEST_CASES  8

int test_case_length[] = {0, 1, 2, 3 , 4, 7, 128, 129};
                                                           
unsigned char keys[] =
{
    0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,
    0x82,0x92,0x5c,0x1c,0xa1,0xd1,0x30,0xb8,
    0x43,0x47,0x21,0xca,0x40,0x63,0x9b,0x3f,
    0xe8,0xf9,0xbe,0xca,0xe9,0x7e,0x5d,0x29,
    0x90,0x03,0x8f,0xc6,0xcf,0x13,0xc1,0xdb,
    0xd5,0x5e,0x10,0x05,0x10,0x12,0x89,0x86
};

unsigned char sas[] =
{

    0x01,0x02,0x03,0x04,0x05,0x06,
    0x01,0x02,0x03,0x04,0x05,0x06,
    0x01,0x02,0x03,0x04,0x05,0x06,
    0x01,0x02,0x03,0x04,0x05,0x06,
    0x01,0x02,0x03,0x04,0x05,0x06,
    0x01,0x02,0x03,0x04,0x05,0x06,
    0x01,0x02,0x03,0x04,0x05,0x06,
    0xab,0x12,0xcd,0x34,0xef,0x56,
};

unsigned char das[] =
{
    0x01,0x22,0x33,0x44,0x55,0x66,
    0x01,0x22,0x33,0x44,0x55,0x66,
    0x01,0x22,0x33,0x44,0x55,0x66,
    0x01,0x22,0x33,0x44,0x55,0x66,
    0x01,0x22,0x33,0x44,0x55,0x66,
    0x01,0x22,0x33,0x44,0x55,0x66,
    0x01,0x22,0x33,0x44,0x55,0x66,
    0x87,0x76,0x65,0x54,0x43,0x32
};

unsigned char priorities[] =
{
    0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x02
};

int do_headers[] =     /* set to zero if da,sa & priority should to be inserted */
{
    0,0,0,0,0,0,1,1
};
unsigned char test_cases[] =
{
    'M',                               /* First 6 test cases come from existing */
    'M','i',                           /* Tgi spec.                             */
    'M','i','c',
    'M','i','c','h',
    'M','i','c','h','a','e','l',

                                /* Test Case 7 */
    0x9a,0xaa,0x03,0x00,0x00,0x00,0x88,0x8e,0x01,0x01,0x00,0x00,0x00,0x00,0x00,0x00,
    0xaa,0xaa,0x03,0x00,0x00,0x00,0x88,0x8e,0x01,0x01,0x00,0x00,0x00,0x00,0x00,0x00,
    0xaa,0xaa,0x03,0x00,0x00,0x00,0x88,0x8e,0x01,0x01,0x00,0x00,0x00,0x00,0x00,0x00,
    0xaa,0xaa,0x03,0x00,0x00,0x00,0x88,0x8e,0x01,0x01,0x00,0x00,0x00,0x00,0x00,0x00,
    0xaa,0xaa,0x03,0x00,0x00,0x00,0x88,0x8e,0x01,0x01,0x00,0x00,0x00,0x00,0x00,0x00,
    0xaa,0xaa,0x03,0x00,0x00,0x00,0x88,0x8e,0x01,0x01,0x00,0x00,0x00,0x00,0x00,0x00,
    0xaa,0xaa,0x03,0x00,0x00,0x00,0x88,0x8e,0x01,0x01,0x00,0x00,0x00,0x00,0x00,0x00,
    0xaa,0xaa,0x03,0x00,0x00,0x00,0x88,0x8e,0x01,0x01,0x00,0x00,0x00,0x00,0x00,0x00,

                                /* Test Case 8 */
    0x8a,0xaa,0x03,0x00,0x00,0x00,0x88,0x8e,0x01,0x01,0x00,0x00,0x00,0x00,0x00,0x00,
    0xaa,0xaa,0x03,0x00,0x00,0x00,0x88,0x8e,0x01,0x01,0x00,0x00,0x00,0x00,0x00,0x00,
    0xaa,0xaa,0x03,0x00,0x00,0x00,0x88,0x8e,0x01,0x01,0x00,0x00,0x00,0x00,0x00,0x00,
    0xaa,0xaa,0x03,0x00,0x00,0x00,0x88,0x8e,0x01,0x01,0x00,0x00,0x00,0x00,0x00,0x00,
    0xaa,0xaa,0x03,0x00,0x00,0x00,0x88,0x8e,0x01,0x01,0x00,0x00,0x00,0x00,0x00,0x00,
    0xaa,0xaa,0x03,0x00,0x00,0x00,0x88,0x8e,0x01,0x01,0x00,0x00,0x00,0x00,0x00,0x00,
    0xaa,0xaa,0x03,0x00,0x00,0x00,0x88,0x8e,0x01,0x01,0x00,0x00,0x00,0x00,0x00,0x00,
    0xaa,0xaa,0x03,0x00,0x00,0x00,0x88,0x8e,0x01,0x01,0x00,0x00,0x00,0x00,0x00,0x00,
    0xba
};



/*****************************/
/**** Function Prototypes ****/
/*****************************/

void get_testcase(  int test_case,
                    unsigned char *plaintext);

int format_message( int do_header,
                    int message_length,
                    unsigned char *message,
                    unsigned char *da,
                    unsigned char *sa,
                    unsigned char priority
                );

void block_function(
                    unsigned long int l,
                    unsigned long int r,
                    unsigned long int *l_out,
                    unsigned long int *r_out
                    );

void michael(
                unsigned char *key,
                unsigned char *message,
                int message_length,
                unsigned char *mic
            );

/**************************/
/* michael()              */
/**************************/

unsigned long int rotr(int bits, unsigned long int a)
{
unsigned long int c,d,e,f,g;
    c = (0x0001 << bits)-1;
    d = ~c;

    e = (a & d) >> bits;
    f = (a & c) << (32 - bits);

    g = e | f;

    return (g & 0xffffffff );

}

unsigned long int rotl(int bits, unsigned long int a)
{
unsigned long int c,d,e,f,g;
    c = (0x0001 << (32-bits))-1;
    d = ~c;

    e = (a & c) << bits;
    f = (a & d) >> (32 - bits);

    g = e | f;

    return (g & 0xffffffff );

}

unsigned long int xswap(unsigned long int in)
{
unsigned long int a,b,c,d,out;

    a = in & 0xff;
    b = (in >> 8) & 0xff;
    c = (in >> 16) & 0xff;
    d = (in >> 24) & 0xff;

    out = 0;
    out |= c << 24;
    out |= d << 16;
    out |= a << 8;
    out |= b;

    return out;
}
void block_function(
                    unsigned long int l,
                    unsigned long int r,
                    unsigned long int *l_out,
                    unsigned long int *r_out
                    )

{

    r = r ^ rotl(17,l);
    l = (l + r) & 0xffffffff;
    r = r ^ xswap(l);
    l = (l + r) & 0xffffffff;
    r = r ^ rotl(3,l);
    l = (l + r) & 0xffffffff;
    r = r ^ rotr(2,l);
    l = (l + r) & 0xffffffff;

    *l_out = l;
    *r_out = r;
}

int format_message( int do_header,
                    int message_length,
                    unsigned char *message,
                    unsigned char *da,
                    unsigned char *sa,
                    unsigned char priority
                )
{
    int length;
    int num_blocks;
    int remainder;

    int i;

    if (do_header)
    {
        length = message_length + 16 + 5; /* da + sa + priority + 000 + 5a + 0000= 21 */
        num_blocks = length / 4;
        remainder = length % 4;

        for (i=0;i<message_length; i++)     /* Shift message right to make room for header */
        {
            message[message_length+15-i] = message[message_length-i-1];
            /* message[i+16] = message[i]; */
        }

        for (i=0; i<6; i++)                 /* Insert addresses */
        {
            message[i] = da[i];
            message[i+6] = sa[i];
        }
        message[12] = priority;    /* Insert priority */
        message[13] = 0x00;
        message[14] = 0x00;
        message[15] = 0x00;

        message[message_length+16] = 0x5a;   /* Insert padding */
        message[message_length+17] = 0x00;
        message[message_length+18] = 0x00;
        message[message_length+19] = 0x00;
        message[message_length+20] = 0x00;
        message[message_length+21] = 0x00;
        message[message_length+22] = 0x00;
        message[message_length+23] = 0x00;

        if (remainder != 0) num_blocks++;
    }
    else
    {
        length = message_length + 5; /* 5a + 0000= 5 */
        num_blocks = length / 4;
        remainder = length % 4;

        message[message_length] = 0x5a;  /* Insert padding */
        message[message_length+1] = 0x00;
        message[message_length+2] = 0x00;
        message[message_length+3] = 0x00;
        message[message_length+4] = 0x00;
        message[message_length+5] = 0x00;
        message[message_length+6] = 0x00;
        message[message_length+7] = 0x00;

        if (remainder != 0) num_blocks++;
        
    }

    return (num_blocks * 4);
}

void michael(
                unsigned char *key,
                unsigned char *message,
                int message_length,
                unsigned char *mic
            )
{
    unsigned long int l,r;
    unsigned long int l_out, r_out;
    unsigned long int k0,k1;
    unsigned long int m;
    unsigned char a,b,c,d;

    int num_blocks;
    int block;

    num_blocks = message_length / 4;

    k0 = key[0] + (key[1] * 256) + (key[2] * 256*256) + (key[3] * 256*256*256);
    k1 = key[4] + (key[5] * 256) + (key[6] * 256*256) + (key[7] * 256*256*256);
    
    l = k0;
    r = k1;

    for (block = 0; block < num_blocks; block++)
    {
        m = 0;
        a = message[block*4];
        b = message[1+(block*4)];
        c = message[2+(block*4)];
        d = message[3+(block*4)];

        m = a + (b * 256) + (c * 256*256) + (d * 256*256*256);
        l = l ^ m;
        block_function(l,r,&l_out,&r_out);
        l = l_out;
        r = r_out;
    }

    mic[0] = (unsigned char)(l & 0xff);
    mic[1] = (unsigned char)((l >> 8) & 0xff);
    mic[2] = (unsigned char)((l >> 16) & 0xff);
    mic[3] = (unsigned char)((l >> 24) & 0xff);
    mic[4] = (unsigned char)(r & 0xff);
    mic[5] = (unsigned char)((r >> 8) & 0xff);
    mic[6] = (unsigned char)((r >> 16) & 0xff);
    mic[7] = (unsigned char)((r >> 24) & 0xff);
}


/************************************************/
/* get_testcase()                               */
/* Copies an mpdu from the test case data       */
/************************************************/
void get_testcase(  int test_case,
                    unsigned char *plaintext)
{
    int i;
    unsigned char *ptr;

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

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

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");
    }
}
/****************************************************/
/* main()                                           */
/* Iterate through the test cases, passing them     */
/* through the ccm algorithm to produce test        */
/* vectors and verify decryption against encryption */
/****************************************************/

int main()
{
    int i;

    int length;
    int test_case;
    int do_header;

    unsigned char plaintext[4096+30];
    unsigned char mic[8];

    unsigned char *key;
    unsigned char *sa;
    unsigned char *da;
    unsigned char priority;

    unsigned long int l,r;
    unsigned long int l_out, r_out;

    printf("Block Function Test\n");
    for (i=0;i<4;i++)
    {
        if (i==0)
        {
            l=0x00000000;
            r=0x00000000;
        }
        else if (i==1)
        {
            l=0x00000000;
            r=0x00000001;
        }
        else if (i==2)
        {
            l=0x00000001;
            r=0x00000000;
        }
        else if (i==3)
        {
            l=0x01234567;
            r=0x83659326;
        }

        block_function(l,r,&l_out,&r_out);
        printf("\t(%08lX)(%08lX)\t1\t(%08lX)(%08lX)\n",l,r,l_out,r_out);
    }

    l=1;
    r=0;
    for (i=0; i<1000;i++)
    {
        block_function(l,r,&l_out,&r_out);
        l= l_out;
        r= r_out;
    }
    printf("\t(%08lX)(%08lX)\t1000\t(%08lX)(%08lX)\n",1UL,0UL,l_out,r_out);

    for (test_case = 1; test_case < (NUM_TEST_CASES+1); test_case++)
    {
        get_testcase(test_case, plaintext);

        key = keys+(8*(test_case-1));
        da = das+(6*(test_case-1));
        sa = sas+(6*(test_case-1));
        priority = priorities[test_case-1];
        do_header = do_headers[test_case-1];
                
        printf ("\nTEST CASE %d\n",test_case);

        debug_out("key =", key, 8);

        if (do_header)
        {
            debug_out("da =", da, 6);
            debug_out("sa =", sa, 6);
            printf("priority = %d\n",(int)priority);
        }
        debug_out("plaintext =", plaintext, test_case_length[test_case-1]);


        length = format_message(do_header, test_case_length[test_case-1], plaintext, da, sa, priority);

        debug_out("formatted_plaintext =", plaintext, length);
        printf("\tformatted_length=%d\n",length);

        michael(key, plaintext, length ,mic);


        debug_out("mic =", mic, 8);
    }

    return 0;
}