/*
 * burt rosenberg
 * feb 2005
 *
 */

#define DEMO

#include<stdio.h>

struct rc4 {
  int i ; 
  int j ; 
  int mask ;
  int * sa ; 
} ;

struct rc4 * rc4_new( int bits ) {
  int mN ;
  struct rc4 * r = (struct rc4 *) malloc( sizeof(struct rc4) ) ;
  mN = 1 << bits ;
  r->mask = mN - 1 ;
  r->sa = (int *) malloc( mN * sizeof(int) ) ;
  return r ;
}

void rc4_init( struct rc4 * r ) {
  int i ;
  r->i = r->j = 0 ;
  for ( i=0; i<=r->mask; i++ ) r->sa[i] = i ;
}

/* rc4_key,
 *    keys the rc4 state
 *    m == 0, do the entire keying operation, including a call to rc4_init
 *    m > 0,  do m rounds of keying, do not rc4_init prior. used to step 
 *            through or for rollig forward to an intermediate stage
 *    m == -1 finish keying. like m for appropriate value to bring keying
 *            to a completion, but sets state i, j to zero for encryption.
 */

void rc4_key( struct rc4 * r, int * key, int k_l, int m ) {
  int mask, m_, t  ;
  m_ = m ;
  mask = r->mask ;

  if ( m==0 ) {
    rc4_init(r) ;
    m = mask + 1 ;
  }

  if ( m<0 ) {
    m = (mask+1) - r->i ;
  }

  while (m--) {
    r->j = (r->j + (t=r->sa[r->i]) + key[r->i%k_l]) &  mask ;
    r->sa[r->i] = r->sa[r->j] ;
    r->sa[r->j] = t ;
    r->i++ ;
  }
  if ( m_<=0 ) r->i = r->j = 0 ;
}


void rc4_crypt( struct rc4 * r, int * buf, int buf_l ) {
  int mask = r->mask ;
  int i = 0, s, t ;

  while ( buf_l-- ) {
    r->i = (r->i + 1) & mask ;
    r->j = (r->j + (t=r->sa[r->i])) & mask  ;
    s = r->sa[r->i] = r->sa[r->j] ;
    r->sa[r->j] = t ;
    buf[i] ^= r->sa[(s+t)&mask] ;
    i++ ;
  }
}

void rc4_print( struct rc4 * r ) {
  int i ;
  printf("\n[%3d %3d] ", r->i, r->j ) ;
  for ( i=0; i<=r->mask; i++ ) printf("%3d ", r->sa[i] ) ;
  printf("\n" ) ;
}

#ifdef TEST

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



/*
 * ARC4 tests vectors from OpenSSL (crypto/rc4/rc4test.c)
 */

static int keys[7][30]={
        {8,0x01,0x23,0x45,0x67,0x89,0xab,0xcd,0xef},
        {8,0x01,0x23,0x45,0x67,0x89,0xab,0xcd,0xef},
        {8,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00},
        {4,0xef,0x01,0x23,0x45},
        {8,0x01,0x23,0x45,0x67,0x89,0xab,0xcd,0xef},
        {4,0xef,0x01,0x23,0x45},
        };

static int data_len[7]={8,8,8,20,28,10};
static int data[7][30]={
        {0x01,0x23,0x45,0x67,0x89,0xab,0xcd,0xef,0xff},
        {0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0xff},
        {0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0xff},
        {0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,
           0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,
           0x00,0x00,0x00,0x00,0xff},
        {0x12,0x34,0x56,0x78,0x9A,0xBC,0xDE,0xF0,
           0x12,0x34,0x56,0x78,0x9A,0xBC,0xDE,0xF0,
           0x12,0x34,0x56,0x78,0x9A,0xBC,0xDE,0xF0,
           0x12,0x34,0x56,0x78,0xff},
        {0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0x00,0xff},
        {0},
        };

static int output[7][30]={
        {0x75,0xb7,0x87,0x80,0x99,0xe0,0xc5,0x96,0x00},
        {0x74,0x94,0xc2,0xe7,0x10,0x4b,0x08,0x79,0x00},
        {0xde,0x18,0x89,0x41,0xa3,0x37,0x5d,0x3a,0x00},
        {0xd6,0xa1,0x41,0xa7,0xec,0x3c,0x38,0xdf,
         0xbd,0x61,0x5a,0x11,0x62,0xe1,0xc7,0xba,
         0x36,0xb6,0x78,0x58,0x00},
        {0x66,0xa0,0x94,0x9f,0x8a,0xf7,0xd6,0x89,
         0x1f,0x7f,0x83,0x2b,0xa8,0x33,0xc0,0x0c,
         0x89,0x2e,0xbe,0x30,0x14,0x3c,0xe2,0x87,
         0x40,0x01,0x1e,0xcf,0x00},
        {0xd6,0xa1,0x41,0xa7,0xec,0x3c,0x38,0xdf,0xbd,0x61,0x00},
        {0},
        };

int main( void )
{
    int i;
    struct rc4 * s;
    int buffer[30];

    printf( "\n RC4 Validation Tests:\n\n" );

    s = rc4_new(8) ;

    for( i = 0; i < 6; i++ )
    {
        printf( " Test %d ", i + 1 );

        memcpy( buffer, data[i], sizeof(int)*data_len[i] );

        rc4_key( s, &keys[i][1], keys[i][0], 0 );
        rc4_crypt( s, buffer, data_len[i] );

        if( memcmp( buffer, output[i], sizeof(int)*data_len[i] ) )
        {
            printf( "failed!\n" );
            return( 1 );
        }

        printf( "passed.\n" );
    }

    printf( "\n" );

    return( 0 );
}

#endif


#ifdef DEMO

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

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

  struct rc4 * r ;
  int i ;
  int key[4] = { 1, 2, 3, 4 } ;
  r = rc4_new(3) ;
  rc4_init( r ) ;
  for ( i=0; i<8; i++ ) {
     rc4_key( r, key, 4, 1 ) ;
     rc4_print( r ) ;
  }
  rc4_key( r, key, 4, -1 ) ;

}

#endif