/*
 *  Copyright (c) 2007-2009 Erik Tews, Andrei Pychkine and Ralf-Philipp
 * Weinmann.
 *                2013 Ramiro Polla
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; either version 2 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program; if not, write to the Free Software
 *  Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301, USA
 *
 *
 *  In addition, as a special exception, the copyright holders give
 *  permission to link the code of portions of this program with the
 *  OpenSSL library under certain conditions as described in each
 *  individual source file, and distribute linked combinations
 *  including the two.
 *  You must obey the GNU General Public License in all respects
 *  for all of the code used other than OpenSSL. *  If you modify
 *  file(s) with this exception, you may extend this exception to your
 *  version of the file(s), but you are not obligated to do so. *  If you
 *  do not wish to do so, delete this exception statement from your
 *  version. *  If you delete this exception statement from all source
 *  files in the program, then also delete it here.
 */

#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include <pthread.h>
#ifdef HAVE_SYS_TYPES_H
#include <sys/types.h>
#endif
#if defined(__sun__)
#include <alloca.h>
#endif
#include <limits.h>
#include <math.h>
#include <float.h>

#include "aircrack-ng/support/pcap_local.h"
#include "aircrack-ng/defs.h"
#include "aircrack-ng/ptw/aircrack-ptw-lib.h"
#include "aircrack-ng/aircrack-ng.h"
#include <aircrack-ng/support/common.h>

#define n PTW_n
#define CONTROLSESSIONS PTW_CONTROLSESSIONS
#define KSBYTES PTW_KSBYTES
#define IVBYTES PTW_IVBYTES
#define TESTBYTES 6

static struct options opt;

// Internal state of rc4
typedef struct
{
	uint32_t s[n];
	uint8_t i;
	uint8_t j;
} rc4state;

// Helper structures for sorting
typedef struct
{
	int keybyte;
	uint8_t value;
	int distance;
} sorthelper;

typedef struct
{
	int keybyte;
	double difference;
} doublesorthelper;

// The rc4 initial state, the idendity permutation
static const uint32_t rc4initial[] = {
	0,   1,   2,   3,   4,   5,   6,   7,   8,   9,   10,  11,  12,  13,  14,
	15,  16,  17,  18,  19,  20,  21,  22,  23,  24,  25,  26,  27,  28,  29,
	30,  31,  32,  33,  34,  35,  36,  37,  38,  39,  40,  41,  42,  43,  44,
	45,  46,  47,  48,  49,  50,  51,  52,  53,  54,  55,  56,  57,  58,  59,
	60,  61,  62,  63,  64,  65,  66,  67,  68,  69,  70,  71,  72,  73,  74,
	75,  76,  77,  78,  79,  80,  81,  82,  83,  84,  85,  86,  87,  88,  89,
	90,  91,  92,  93,  94,  95,  96,  97,  98,  99,  100, 101, 102, 103, 104,
	105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119,
	120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134,
	135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149,
	150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164,
	165, 166, 167, 168, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179,
	180, 181, 182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 193, 194,
	195, 196, 197, 198, 199, 200, 201, 202, 203, 204, 205, 206, 207, 208, 209,
	210, 211, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 222, 223, 224,
	225, 226, 227, 228, 229, 230, 231, 232, 233, 234, 235, 236, 237, 238, 239,
	240, 241, 242, 243, 244, 245, 246, 247, 248, 249, 250, 251, 252, 253, 254,
	255};

// Values for p_correct_i
static const double eval[] = {0.00534392069257663,
							  0.00531787585068872,
							  0.00531345769225911,
							  0.00528812219217898,
							  0.00525997750378221,
							  0.00522647312237696,
							  0.00519132541143668,
							  0.0051477139367225,
							  0.00510438884847959,
							  0.00505484662057323,
							  0.00500502783556246,
							  0.00495094196451801,
							  0.0048983441590402};

static int tried, max_tries;
static int depth[KEYHSBYTES];
static PTW_tableentry keytable[KEYHSBYTES][n];

// Check if optmizied RC4 for AMD64 has to be compiled
#if defined(__amd64) && defined(__SSE2__)                                      \
	&& (!defined(__clang__)                                                    \
		|| (defined(__clang__)                                                 \
			&& (__clang_major__ >= 4                                           \
				|| (__clang_major__ == 3 && __clang_minor__ >= 9))))
#define USE_AMD64_RC4_OPTIMIZED
#endif

// For sorting
static int compare(const void * ina, const void * inb)
{
	REQUIRE(ina != NULL);
	REQUIRE(inb != NULL);

	PTW_tableentry * a = (PTW_tableentry *) ina;
	PTW_tableentry * b = (PTW_tableentry *) inb;
	return b->votes - a->votes;
}

// For sorting
static int comparedoublesorthelper(const void * ina, const void * inb)
{
	REQUIRE(ina != NULL);
	REQUIRE(inb != NULL);

	doublesorthelper * a = (doublesorthelper *) ina;
	doublesorthelper * b = (doublesorthelper *) inb;
	if (a->difference > b->difference)
	{
		return 1;
	}
	else if (fabs(a->difference - b->difference) < FLT_EPSILON)
	{
		return 0;
	}
	else
	{
		return -1;
	}
}

#ifdef USE_AMD64_RC4_OPTIMIZED
static const uint32_t __attribute__((used)) __attribute__((aligned(16)))
x0123[4]
	= {0, 1, 2, 3};
static const uint32_t __attribute__((used)) __attribute__((aligned(16)))
x4444[4]
	= {4, 4, 4, 4};
static int
rc4test_amd64_sse2(uint8_t * key, int keylen, uint8_t * iv, uint8_t * keystream)
{
	int idx, i, j;
	int scratch1, scratch2;

	__asm__ volatile(
#define state "%%rsp"
#define keybuf "0x400(%%rsp)"
#define keystream_ "0x428(%%rsp)"
		// setup stack
		"movq  %%rsp, %q0             \n\t"
		"subq $0x430, %%rsp           \n\t"
		"andq   $-16, %%rsp           \n\t"
		"movq    %q0, -8(%%rsp)       \n\t"

		// save keystream variable
		"movq %q6, " keystream_ "       \n\t"

		// keylen += IVBYTES
		"addl    $3, %k4              \n\t"

		// memcpy(keybuf, iv, IVBYTES);
		"movl  (%q5), %k1             \n\t"
		"movl   %k1 , " keybuf "        \n\t"
		// memcpy(&keybuf[IVBYTES], key, keylen);
		"movdqa   (%q3), %%xmm0       \n\t"
		"cmpl    $16, %k4             \n\t"
		"movdqu %%xmm0, 3+" keybuf "    \n\t"
		"jng     .Lsmall_key1         \n\t"
		"movdqa 16(%q3), %%xmm1       \n\t"
		"movdqu %%xmm1,19+" keybuf "    \n\t"
		".Lsmall_key1:                \n\t"

		// key = keybuf
		"lea  " keybuf ", %q3           \n\t"
		// load xmm registers
		"movdqa %q9,  %%xmm0          \n\t"
		"movdqa %q10, %%xmm1          \n\t"
		// clear some registers
		"xorq      %q0, %q0           \n\t" // idx
		"xorq      %q1, %q1           \n\t" // i
		"xorq      %q2, %q2           \n\t" // j

		// build identity array
		".p2align 4                   \n\t"
		".Lidentity_loop:             \n\t"
		"movdqa %%xmm0, (" state ",%q1,4)\n\t"
		"addb   $4, %b1               \n\t"
		"paddd  %%xmm1, %%xmm0        \n\t"
		"jnc  .Lidentity_loop         \n\t"

		// load state into register
		"movq " state ", %q1            \n\t"

		// %q4 = and mask for idx
		"movq %q4, %q8                \n\t"
		"cmpq $16, %q8                \n\t"
		"movq $15, %q4                \n\t"
		"je    .Lsmall_key2           \n\t"
		"shrq  $1, %q4                \n\t"
		".Lsmall_key2:                \n\t"

		// init array with key
		".p2align 4                   \n\t"
		".init_loop:                  \n\t"
		"movl    %k0, %k8             \n\t" /* scratch2        = idx */
		"movl   (%q1), %k5            \n\t" /* s1              = state[i] */
		"leal  1(%q0,1), %k0          \n\t" /* idx++ */
		"movzbl (%q3,%q8,1), %k6      \n\t" /* key_n           = key[scratch2]
											   */
		"leal   (%q5,%q6,1), %k8      \n\t" /* scratch2        = s1 + key_n */
		"addl    %k8, %k2             \n\t" /* j              += scratch2 */
		"andl    %k4, %k0             \n\t" /* idx            &= mask */
		"movzbl  %b2, %k8             \n\t" /* scratch2        = j */
		"movl (" state
		",%q8,4), %k7    \n\t" /* s2              = state[scratch2] */
		"movl    %k7, (%q1)           \n\t" /* state[i]        = s2 */
		"addq     $4, %q1             \n\t" /* i++ */
		"movl    %k5, (" state ",%q8,4) \n\t" /* state[scratch2] = s1 */
		"cmpq    %q1, %q3             \n\t" /* state          == &state[0x100]
												 */
		"jne .init_loop               \n\t"

		// restore keystream variable
		"movq " keystream_ ", %q6       \n\t"

		// clear some registers
		"xorq  %q2, %q2               \n\t" // j = 0
		"xorq  %q0, %q0               \n\t" // result

#define RC4TEST_LOOP(offset)                                                   \
	"movl 4*" offset "(" state "), %k5\n\t" /* s1 = state[i]         */        \
	"leal (%q5,%q2,1), %k4        \n\t" /*                       */            \
	"movzbl %b4, %k2              \n\t" /* j += s1               */            \
	"movl (" state ",%q2,4), %k1    \n\t" /* s2 = state[j]         */          \
	"movl %k1, 4*" offset "(" state ")\n\t" /* state[i] = s2         */        \
	"movl %k5, (" state ",%q2,4)    \n\t" /* state[j] = s1         */          \
	"addb %b1, %b5                \n\t" /* s1 += s2;             */            \
	"movb (" state ",%q5,4), %b3    \n\t" /* ret = state[s1]       */          \
	"cmpb %b3, " offset "-1(%q6)    \n\t" /* ret == keystream[i-1] */          \
	"jne .ret                     \n\t"

		RC4TEST_LOOP("1") RC4TEST_LOOP("2") RC4TEST_LOOP("3") RC4TEST_LOOP("4")
			RC4TEST_LOOP("5") RC4TEST_LOOP("6")

#undef RC4TEST_LOOP

				"addb $1, %b0                 \n\t"
				".ret:                        \n\t"

				// restore stack
				"movq -8(%%rsp), %%rsp        \n\t"

		: "=&r"(idx),
		  "=&r"(i),
		  "=&r"(j),
		  "+r"(key),
		  "+r"(keylen),
		  "+r"(iv),
		  "+r"(keystream),
		  "=&r"(scratch1),
		  "=&r"(scratch2)
		: "m"(x0123[0]), "m"(x4444[0])
		: "xmm0", "xmm1");
#undef state
#undef keybuf
#undef keystream_

	return idx;
}
#endif

// RC4 key setup
static void rc4init(uint8_t * key, int keylen, rc4state * state)
{
	REQUIRE(key != NULL);
	REQUIRE(keylen > 0 && keylen < INT_MAX);
	REQUIRE(state != NULL);

	int i;
	unsigned char j;
	uint8_t tmp;
	int idx = 0;
	memcpy(state->s, &rc4initial, sizeof(rc4initial));
	j = 0;
	for (i = 0; i < n; i++)
	{
		/*  this should be:
			j = (j + state->s[i] + key[i % keylen]) % n;
			but as "j" is declared as unsigned char and n equals 256,
			we can "optimize" it
		*/
		j = (j + state->s[i] + key[idx]);
		if (++idx == keylen) idx = 0;
		tmp = state->s[i];
		state->s[i] = state->s[j];
		state->s[j] = tmp;
	}
	state->i = 0;
	state->j = 0;
}

// RC4 key stream generation
static uint8_t rc4update(rc4state * state)
{
	REQUIRE(state != NULL);

	uint8_t tmp;
	uint8_t k;
	state->i++;
	state->j += state->s[state->i];
	tmp = state->s[state->i];
	state->s[state->i] = state->s[state->j];
	state->s[state->j] = tmp;
	k = state->s[state->i] + state->s[state->j];

	return state->s[k];
}

static int rc4test(uint8_t * key, int keylen, uint8_t * iv, uint8_t * keystream)
{
	REQUIRE(key != NULL);
	REQUIRE(keylen > 0);
	REQUIRE(iv != NULL);
	REQUIRE(keystream != NULL);

	uint8_t keybuf[PTW_KSBYTES];
	rc4state rc4state;
	int j;
	memcpy(&keybuf[IVBYTES], key, keylen);
	memcpy(keybuf, iv, IVBYTES); //-V512
	rc4init(keybuf, keylen + IVBYTES, &rc4state);
	for (j = 0; j < TESTBYTES; j++)
	{
		if ((rc4update(&rc4state) ^ keystream[j]) != 0)
		{
			return 0;
		}
	}
	return 1;
}

// For sorting
static int comparesorthelper(const void * ina, const void * inb)
{
	REQUIRE(ina != NULL);
	REQUIRE(inb != NULL);

	sorthelper * a = (sorthelper *) ina;
	sorthelper * b = (sorthelper *) inb;
	return a->distance - b->distance;
}

/*
 * Guess the values for sigma_i
 * ivlen - how long was the iv (is used differently in original klein attack)
 * iv - IV which was used for this packet
 * keystream - keystream recovered
 * result - buffer for the values of sigma_i
 * kb - how many keybytes should be guessed
 */
static void guesskeybytes(
	int ivlen, uint8_t * iv, uint8_t * keystream, uint8_t * result, int kb)
{
	REQUIRE(iv != NULL);
	REQUIRE(keystream != NULL);
	REQUIRE(result != NULL);

	uint32_t state[n];
	uint8_t j = 0;
	uint8_t tmp;
	int i;
	int jj = ivlen;
	uint8_t ii;
	uint8_t s = 0;
	memcpy(state, &rc4initial, sizeof(rc4initial));
	for (i = 0; i < ivlen; i++)
	{
		j += state[i] + iv[i];
		tmp = state[i];
		state[i] = state[j];
		state[j] = tmp;
	}
	for (i = 0; i < kb; i++)
	{
		tmp = jj - keystream[jj - 1];
		ii = 0;
		while (tmp != state[ii])
		{
			ii++;
		}
		s += state[jj];
		ii -= (j + s);
		result[i] = ii;
		jj++;
	}
	return;
}

#define PTW_CORRECTNESS_MIN_SESSION_COUNT 3
#define PTW_CORRECTNESS_SESSION_COUNT 10
/*
 * Is a guessed key correct?
 */
static int correct(PTW_attackstate * state, uint8_t * key, int keylen)
{
	REQUIRE(state != NULL);
	REQUIRE(key != NULL && keylen > 0);

	int i;
	int end_check;

	// We need at least 3 sessions to be somehow certain
	if (state->sessions_collected < PTW_CORRECTNESS_MIN_SESSION_COUNT)
	{
		return 0;
	}

	tried++;

	// Validate against 10 sessions at a random place ...
	if (state->sessions_collected > PTW_CORRECTNESS_SESSION_COUNT)
	{
		i = rand_u32() % (state->sessions_collected - PTW_CORRECTNESS_SESSION_COUNT + 1);
		end_check = i + PTW_CORRECTNESS_SESSION_COUNT;
	}
	else
	{
		// ... or against all of them if we don't have 10
		i = 0;
		end_check = state->sessions_collected;
	}
	for (; i < end_check; i++)
	{
		if (!state->rc4test(key,
							keylen,
							state->sessions[i].iv,
							state->sessions[i].keystream))
			return 0;
	}
	return 1;
}

/*
 * Calculate the squaresum of the errors for both distributions
 */
static void getdrv(PTW_tableentry orgtable[][n],
				   int keylen,
				   double * normal,
				   double * ausreiser)
{
	int i, j;
	int numvotes = 0;
	double e;
	double e2;
	double emax;
	double help = 0.0;
	double maxhelp = 0;
	double maxi = 0;
	for (i = 0; i < n; i++)
	{
		numvotes += orgtable[0][i].votes;
	}
	e = (double) (numvotes) / n;
	for (i = 0; i < keylen; i++)
	{
		emax = eval[i] * numvotes;
		e2 = ((1.0 - eval[i]) / 255.0) * numvotes;
		normal[i] = 0;
		ausreiser[i] = 0;
		maxhelp = 0;
		maxi = 0;
		for (j = 0; j < n; j++)
		{
			if (orgtable[i][j].votes > maxhelp)
			{
				maxhelp = orgtable[i][j].votes;
				maxi = j;
			}
		}
		for (j = 0; j < n; j++)
		{
			if (fabs(maxi - j) < FLT_EPSILON)
			{
				help = (1.0 - orgtable[i][j].votes / emax);
			}
			else
			{
				help = (1.0 - orgtable[i][j].votes / e2);
			}
			help = help * help;
			ausreiser[i] += help;
			help = (1.0 - orgtable[i][j].votes / e);
			help = help * help;
			normal[i] += help;
		}
	}
}

/*
 * Guess a single keybyte
 */
static int doRound(PTW_tableentry sortedtable[][n],
				   int keybyte,
				   int fixat,
				   uint8_t fixvalue,
				   int * searchborders,
				   uint8_t * key,
				   int keylen,
				   PTW_attackstate * state,
				   uint8_t sum,
				   int * strongbytes,
				   int * bf,
				   int validchars[][n])
{
	int i;
	uint8_t tmp;

	if (!opt.is_quiet && keybyte < 4)
		show_wep_stats(keylen - 1, 0, keytable, searchborders, depth, tried);
	if (keybyte > 0)
	{
		if (!validchars[keybyte - 1][key[keybyte - 1]])
		{
			return 0;
		}
	}
	if (keybyte == keylen)
	{
		return correct(state, key, keylen);
	}
	else if (bf[keybyte] == 1)
	{
		for (i = 0; i < n; i++)
		{
			key[keybyte] = i;
			if (doRound(sortedtable,
						keybyte + 1,
						fixat,
						fixvalue,
						searchborders,
						key,
						keylen,
						state,
						sum + i % n,
						strongbytes,
						bf,
						validchars))
			{
				return 1;
			}
		}
		return 0;
	}
	else if (keybyte == fixat)
	{
		key[keybyte] = fixvalue - sum;
		return doRound(sortedtable,
					   keybyte + 1,
					   fixat,
					   fixvalue,
					   searchborders,
					   key,
					   keylen,
					   state,
					   fixvalue,
					   strongbytes,
					   bf,
					   validchars);
	}
	else if (strongbytes[keybyte] == 1)
	{
		// printf("assuming byte %d to be strong\n", keybyte);
		tmp = 3 + keybyte;
		for (i = keybyte - 1; i >= 1; i--)
		{
			tmp += 3 + key[i] + i;
			key[keybyte] = n - tmp;
			if (doRound(sortedtable,
						keybyte + 1,
						fixat,
						fixvalue,
						searchborders,
						key,
						keylen,
						state,
						(n - tmp + sum) % n,
						strongbytes,
						bf,
						validchars)
				== 1)
			{
				printf("hit with strongbyte for keybyte %d\n", keybyte);
				return 1;
			}
		}
		return 0;
	}
	else
	{
		REQUIRE(searchborders != NULL);
		REQUIRE(keybyte >= 0);

		for (i = 0; i < searchborders[keybyte]; i++)
		{
			key[keybyte] = sortedtable[keybyte][i].b - sum;
			if (!opt.is_quiet)
			{
				depth[keybyte] = i;
				keytable[keybyte][i].b = key[keybyte];
			}
			if (doRound(sortedtable,
						keybyte + 1,
						fixat,
						fixvalue,
						searchborders,
						key,
						keylen,
						state,
						sortedtable[keybyte][i].b,
						strongbytes,
						bf,
						validchars))
			{
				return 1;
			}
		}
		return 0;
	}
}

/*
 * Do the actual computation of the key
 */
static int doComputation(PTW_attackstate * state,
						 uint8_t * key,
						 int keylen,
						 PTW_tableentry table[][n],
						 sorthelper * sh2,
						 int * strongbytes,
						 int keylimit,
						 int * bf,
						 int validchars[][n])
{
	int i, j;
	int choices[KEYHSBYTES];
	int prod;
	int fixat;
	int fixvalue;

	if (!opt.is_quiet)
		memcpy(keytable, table, sizeof(PTW_tableentry) * n * keylen);

	for (i = 0; i < keylen; i++)
	{
		if (strongbytes[i] == 1)
		{
			choices[i] = i;
		}
		else
		{
			choices[i] = 1;
		}
	}
	i = 0;
	prod = 0;
	fixat = -1;
	fixvalue = 0;
	max_tries = keylimit;

	while (prod < keylimit)
	{
		if (doRound(table,
					0,
					fixat,
					fixvalue,
					choices,
					key,
					keylen,
					state,
					0,
					strongbytes,
					bf,
					validchars)
			== 1)
		{
			// printf("hit with %d choices\n", prod);
			if (!opt.is_quiet)
				show_wep_stats(keylen - 1, 1, keytable, choices, depth, tried);
			return 1;
		}
		while ((i < keylen * (n - 1)) && ((strongbytes[sh2[i].keybyte] == 1)
										  || (bf[sh2[i].keybyte] == 1)))
		{
			i++;
		}
		if (i >= (keylen * (n - 1)))
		{
			break;
		}
		choices[sh2[i].keybyte]++;
		fixat = sh2[i].keybyte;
		// printf("choices[%d] is now %d\n", sh2[i].keybyte,
		// choices[sh2[i].keybyte]);
		fixvalue = sh2[i].value;
		prod = 1;
		for (j = 0; j < keylen; j++)
		{
			prod *= choices[j];
			if (bf[j] == 1)
			{
				prod *= n;
			}
		}

		/*
		do {
			i++;
		} while (strongbytes[sh2[i].keybyte] == 1);
		*/
		i++;

		if (!opt.is_quiet)
			show_wep_stats(keylen - 1, 0, keytable, choices, depth, tried);
	}
	if (!opt.is_quiet)
		show_wep_stats(keylen - 1, 1, keytable, choices, depth, tried);
	return 0;
}

/*
 * Guess which key bytes could be strong and start actual computation of the key
 */
int PTW_computeKey(PTW_attackstate * state,
				   uint8_t * keybuf,
				   int keylen,
				   int testlimit,
				   int * bf,
				   int validchars[][n],
				   int attacks)
{
	REQUIRE(state != NULL);

	int strongbytes[KEYHSBYTES];
	double normal[KEYHSBYTES];
	double ausreisser[KEYHSBYTES];
	doublesorthelper helper[KEYHSBYTES];
	int simple, onestrong, twostrong;
	int i, j;
#ifdef USE_AMD64_RC4_OPTIMIZED
	/*
	 * The 64-bit SSE2-optimized rc4test() requires this buffer to be
	 * aligned at 3 bytes.
	 */
	uint8_t fullkeybuf_unaligned[PTW_KSBYTES + 13] __attribute__((aligned(16)));
	uint8_t * fullkeybuf = &fullkeybuf_unaligned[13];
#else
	uint8_t fullkeybuf[PTW_KSBYTES];
#endif
	uint8_t guessbuf[PTW_KSBYTES];
	sorthelper(*sh)[n - 1];
	PTW_tableentry(*table)[n] = alloca(sizeof(PTW_tableentry) * n * keylen);
	ALLEGE(table != NULL);

#ifdef USE_AMD64_RC4_OPTIMIZED
	/*
	 * sse2-optimized rc4test() function for amd64 only works
	 * for keylen == 5 or keylen == 13
	 */
	if (keylen == 5 || keylen == 13)
		state->rc4test = rc4test_amd64_sse2;
	else
#endif
		state->rc4test = rc4test;

	tried = 0;
	sh = NULL;

	if (!(attacks & NO_KLEIN))
	{
		// Try the original klein attack first
		for (i = 0; i < keylen; i++)
		{
			memset(&table[i][0], 0, sizeof(PTW_tableentry) * n);
			for (j = 0; j < n; j++)
			{
				table[i][j].b = j;
			}
			for (j = 0; j < state->packets_collected; j++)
			{
				// fullkeybuf[0] = state->allsessions[j].iv[0];
				memcpy(
					fullkeybuf, state->allsessions[j].iv, 3 * sizeof(uint8_t));
				guesskeybytes(i + 3,
							  fullkeybuf,
							  state->allsessions[j].keystream,
							  guessbuf,
							  1);
				table[i][guessbuf[0]].votes += state->allsessions[j].weight;
			}
			qsort(&table[i][0], n, sizeof(PTW_tableentry), &compare);
			j = 0;
			while (!validchars[i][table[i][j].b])
			{
				j++;
			}
			// printf("guessing i = %d, b = %d\n", i, table[0][0].b);
			fullkeybuf[i + 3] = table[i][j].b;
		}
		if (correct(state, &fullkeybuf[3], keylen))
		{
			memcpy(keybuf, &fullkeybuf[3], keylen * sizeof(uint8_t));
			// printf("hit without correction\n");
			return (FAILURE);
		}
	}

	if (!(attacks & NO_PTW))
	{
		memcpy(table, state->table, sizeof(PTW_tableentry) * n * keylen);

		onestrong = (testlimit / 10) * 2;
		twostrong = (testlimit / 10) * 1;
		simple = testlimit - onestrong - twostrong;

		// now, sort the table
		for (i = 0; i < keylen; i++)
		{
			qsort(&table[i][0], n, sizeof(PTW_tableentry), &compare);
			strongbytes[i] = 0;
		}

		sh = alloca(sizeof(sorthelper) * (n - 1) * keylen);
		ALLEGE(sh != NULL);

		for (i = 0; i < keylen; i++)
		{
			for (j = 1; j < n; j++)
			{
				sh[i][j - 1].distance = table[i][0].votes - table[i][j].votes;
				sh[i][j - 1].value = table[i][j].b;
				sh[i][j - 1].keybyte = i;
			}
		}
		qsort(sh, (n - 1) * keylen, sizeof(sorthelper), &comparesorthelper);

		if (doComputation(state,
						  keybuf,
						  keylen,
						  table,
						  (sorthelper *) sh,
						  strongbytes,
						  simple,
						  bf,
						  validchars))
		{
			return (FAILURE);
		}

		// Now one strong byte
		getdrv(state->table, keylen, normal, ausreisser);
		for (i = 0; i < keylen - 1; i++)
		{
			helper[i].keybyte = i + 1;
			helper[i].difference = normal[i + 1] - ausreisser[i + 1];
		}
		qsort(helper,
			  keylen - 1,
			  sizeof(doublesorthelper),
			  &comparedoublesorthelper);
		// do not use bf-bytes as strongbytes
		i = 0;
		while (bf[helper[i].keybyte] == 1)
		{
			i++;
		}
		strongbytes[helper[i].keybyte] = 1;
		if (doComputation(state,
						  keybuf,
						  keylen,
						  table,
						  (sorthelper *) sh,
						  strongbytes,
						  onestrong,
						  bf,
						  validchars))
		{
			return (FAILURE);
		}

		// two strong bytes
		i++;
		while (bf[helper[i].keybyte] == 1)
		{
			i++;
		}
		strongbytes[helper[i].keybyte] = 1;
		if (doComputation(state,
						  keybuf,
						  keylen,
						  table,
						  (sorthelper *) sh,
						  strongbytes,
						  twostrong,
						  bf,
						  validchars))
		{
			return (FAILURE);
		}
	}
	return (SUCCESS);
}

/*
 * Add a new session to the attack
 * state - state of attack
 * iv - IV used in the session
 * keystream - recovered keystream from the session
 */
int PTW_addsession(PTW_attackstate * state,
				   uint8_t * iv,
				   uint8_t * keystream,
				   int * weight,
				   int total)
{
	REQUIRE(state != NULL);
	REQUIRE(iv != NULL);
	REQUIRE(keystream != NULL);
	REQUIRE(weight != NULL);

	int i, j;
	int il;
	int ir;
	uint8_t buf[PTW_KEYHSBYTES];

	i = (iv[0] << 16) | (iv[1] << 8) | (iv[2]);
	il = i / 8;
	ir = 1 << (i % 8);
	if ((state->seen_iv[il] & ir) == 0)
	{
		state->seen_iv[il] |= ir;
		for (j = 0; j < total; j++)
		{
			state->packets_collected++;
			guesskeybytes(
				IVBYTES, iv, &keystream[KSBYTES * j], buf, PTW_KEYHSBYTES);
			for (i = 0; i < KEYHSBYTES; i++)
			{
				state->table[i][buf[i]].votes += weight[j];
			}
			if (state->allsessions_size < state->packets_collected)
			{
				state->allsessions_size = state->allsessions_size << 1;
				PTW_session * tmp_allsessions
					= realloc(state->allsessions,
							  state->allsessions_size * sizeof(PTW_session));
				ALLEGE(tmp_allsessions != NULL);
				state->allsessions = tmp_allsessions;
			}
			memcpy(state->allsessions[state->packets_collected - 1].iv,
				   iv,
				   IVBYTES);
			memcpy(state->allsessions[state->packets_collected - 1].keystream,
				   &keystream[KSBYTES * j],
				   KSBYTES);
			state->allsessions[state->packets_collected - 1].weight = weight[j];
		}
		if ((state->sessions_collected < CONTROLSESSIONS))
		{
			memcpy(state->sessions[state->sessions_collected].iv, iv, IVBYTES);
			memcpy(state->sessions[state->sessions_collected].keystream,
				   keystream,
				   KSBYTES);
			state->sessions_collected++;
		}

		return (FAILURE);
	}
	else
	{
		return (SUCCESS);
	}
}

/*
 * Allocate a new attackstate
 */
PTW_attackstate * PTW_newattackstate(void)
{
	int i, k;
	PTW_attackstate * state = NULL;
	state = malloc(sizeof(PTW_attackstate));
	ALLEGE(state != NULL);
	memset(state, 0, sizeof(PTW_attackstate));
	for (i = 0; i < PTW_KEYHSBYTES; i++)
	{
		for (k = 0; k < n; k++)
		{
			state->table[i][k].b = k;
		}
	}
	state->allsessions = malloc(4096 * sizeof(PTW_session));
	ALLEGE(state->allsessions != NULL);
	state->allsessions_size = 4096;

	return state;
}

/*
 * Free an allocated attackstate
 */
void PTW_freeattackstate(PTW_attackstate * state)
{
	free(state->allsessions);
	free(state);
	return;
}
