/*
------------------------------------------------------------------------
Bob Jenkins, 1996

Code for breaking rc4.  Doesn't work.
------------------------------------------------------------------------
*/

#define ALPHA 6
#define SIZE  (1<<ALPHA)
#define MASK  (SIZE-1)
#define RSIZE (SIZE*2)
#define MYRUNS 100000

/*
------------------------------------------------------------------------
rc4 and rc4b
------------------------------------------------------------------------
*/

struct rc4state
{
  int  j;        /* accumulator */
  int  m[SIZE];  /* memory */
  int  r[SIZE];  /* last SIZE results */
};
typedef  struct rc4state  rc4state;

/* RC4 forward SIZE steps.  Sets s->r correctly. */
static void rc4( rc4state *s)
{
  int i,j,x,y,*m,*r;

  j = s->j;
  m = s->m;
  r = s->r;
  for (i=0; i<SIZE; ++i)
  {
    x = m[i];
    j = (x+j) & MASK;
    y = m[j];
    m[i] = y;
    m[j] = x;
    r[i] = m[((x+y)&MASK)];
  }
  s->j = j;
}

/* RC4 backwards SIZE steps.  Leaves s->r invalid. */
static void rc4b( rc4state *s)
{
  int i,j,x,y,*m;

  j = s->j;
  m = s->m;
  for (i=SIZE; i--; )
  {
    x=m[i];
    y=m[j];
    m[i]=y;
    m[j]=x;
    j=(j-y)&MASK;
  }
  s->j = j;
}

/* Copy the state of rc4 */
static void rc4copy( rc4state *outs, rc4state *ins)
{
  int i;
  outs->j = ins->j;
  for (i=0; i<SIZE; ++i)
  {
    outs->m[i] = ins->m[i];
    outs->r[i] = ins->r[i];
  }
}

/* print out a message if s1 and s2 are not the same */
static int rc4test( rc4state *s1, rc4state *s2)
{
  int i;
  if (s1->j != s2->j)
    printf( "rc4test: Accumulator is different, %d %d\n", s1->j, s2->j);
  for (i=0; i<SIZE; ++i)
    if (s1->m[i] != s2->m[i])
      printf( "rc4test: memory [%d] is different, %d %d\n", 
              i, s1->m[i], s2->m[i]);
}

/* initialize rc4 */
int rc4init( rc4state *s)
{
  int i;
  int r[SIZE];

  s->j = 10;
  for (i=0; i<SIZE; ++i) 
    s->m[i]=i;
}

/*
------------------------------------------------------------------------
Source of random values
------------------------------------------------------------------------
*/

static int      rng_i;
static rc4state rng_s;

int rng()                                         /* return one random value */
{
  if (++rng_i >= SIZE)
  {
    rc4( &rng_s);
    rng_i=0;
  }
  return rng_s.r[rng_i];
}


/*
------------------------------------------------------------------------
Tools for breaking rc4
------------------------------------------------------------------------
*/

/*
 * Count how many relevant r values match the stream generated by s.
 *
 * s - rc4 state, already run.
 * n - array size
 * r - array of results to compare against.
 * mask - array same size as r, true/false, pay attention to this value
 */
static int count_match( rc4state *s, int n, int *r, int *mask)
{
  int  i, count;

  count = 0;

  if (n <= SIZE)
  {
    for (i=0; i<n; ++i)
      if (mask[i] && (s->r[i] == r[i]))
        ++count;
  }
  else
  {
    rc4state s2;
    int      j;

    rc4copy( &s2, s);

    for (j=0; j<n-SIZE; j+=SIZE)
    {
      for (i=0; i<SIZE; ++i)
        if (mask[i+j] && (s2.r[i] == r[i+j]))
          ++count;
      rc4( &s2);
    }

    for (i=0; j+i<n; ++i)
      if (mask[j+i] && (s2.r[i] == r[j+i]))
        ++count;
  }

  return count;
}



/*
 * Given an already-initialized s, try to maximize relevant matches to r
 *
 * n - array size
 * r - array of results to match
 * mask - array of true/false, whether that entry in r is relevant
 * m - stop once you get this many matches
 */
static int climb_hill( rc4state *s, int n, int *r, int *mask, int m)
{
  int i;
  int best;                            /* best number of matches seen so far */
  int this;                               /* number of matches in this guess */

  best = count_match( s, n, r, mask);

  for (i=0; i<MYRUNS; ++i)
  {
    int j;
    int x,y;       /* x and y are the random values, z is used to build them */

    for (;;) 
    {                                          /* get distinct random values */
      x=rng();
      y=rng();
      if (x!=y) break;
    }
    
    {j=s->m[x]; s->m[x]=s->m[y]; s->m[y]=j;}                         /* swap */
    rc4b( s);
    rc4( s);

    this = count_match( s, n, r, mask);

    if (this>best)  
    {                                        /* improvement, report progress */
      /* printf( "this is %4d\n",this); */
      i = 0;
      best = this;
    }
    else if (this==best) ;             /* not losing ground, keep the change */
    else {j=s->m[x]; s->m[x]=s->m[y]; s->m[y]=j;}          /* bad, so unswap */

    if (this==m)
      return 1;                                   /* found a solution?  stop */
  }
  return 0;
}

int main()
{
  long int i, j;                                                  /* counter */
  int k = 10;                                   /* how many relevant entries */
  int r[RSIZE];                                          /* results to match */
  int maska[RSIZE];                            /* which results are relevant */
  int mask[RSIZE];                             /* which results are relevant */
  int hit[SIZE];
  rc4state s;                                          /* rc4 internal state */
  rc4state s2;                                       /* extra internal state */

  /* initialize randomness source */
  rc4init( &rng_s);
  rc4( &rng_s);
  rc4( &rng_s);
  rc4( &rng_s);
  rc4( &rng_s);
  rc4( &rng_s);
  rng_i = 0;

  /* confirm that rc4 and rc4b work correctly, and initialize s->r */
  rc4init( &s);
  rc4copy( &s2, &s);
  rc4b( &s);
  rc4( &s);
  rc4test( &s, &s2);

  /* initialize the results to be matched */
  for (i=0; i<RSIZE; ++i)
  {
    r[i]=((i/2)&MASK);
    rng();

    if (i > SIZE-k && i < SIZE+k && i != SIZE)
      mask[i] = 1;
    else
      mask[i] = 0;

    if (i > SIZE-k/2 && i < SIZE+k/2 && i != SIZE)
      maska[i] = 1;
    else
      maska[i] = 0;
  }

  /* no hits so far */
  for (i=0; i<SIZE; ++i) 
    hit[i] = 0;

  /* give it a shot */
  for (i=0;;)
  {
    /* try to guess the state that produces those results */
    j = climb_hill( &s, RSIZE, r, maska, 2*(k/2)-2);
    j = climb_hill( &s, RSIZE, r, mask, 2*k-2);
    rc4(&s);
    if (j) ++hit[s.r[0]];
    if (!j) continue;
    ++i;
    for (j=0; j<SIZE; ++j) 
    {
      printf("%6.2f", SIZE*(hit[j]/(double)i));
      if ((j%12) == 11) printf("\n");
    }
    printf("\n %ld \n", i);
  }
}