/* rc4tiny.c
 * slow and simple RC4 implementation mainly for random number generation
 *
 * Jon Mayo - PUBLIC DOMAIN - January 11, 2011
 */
#include <stdint.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <inttypes.h>

struct rc4_key;
void rc4_set_key(struct rc4_key *key, size_t len, const void *data);
void rc4(struct rc4_key *key, size_t len, const void *in, void *out);
uint32_t rc4rand32(struct rc4_key *key);
uint64_t rc4rand64(struct rc4_key *key);

struct rc4_key
{
	uint_fast8_t x, y;
	uint_fast8_t data[256];
};

static inline void xchg(struct rc4_key *key, unsigned x, unsigned y)
{
	uint_fast8_t tmp = key->data[x];
	key->data[x] = key->data[y];
	key->data[y] = tmp;
}

/* initialize with a key of any size */
void rc4_set_key(struct rc4_key *key, size_t len, const void *data)
{
	unsigned i, j;
	for (i = 0; i < 256; i++)
		key->data[i] = i;

	for (i = j = 0; i < 256; i++) {
		j = (j + ((unsigned char *)data)[i % len] + key->data[i]) & 255;
		xchg(key, i, j);
	}
	key->x = key->y = 0;
}

/* encrypt */
void rc4(struct rc4_key *key, size_t len, const void *in, void *out)
{
	unsigned i;
	for (i = 0; i < len; i++) {
		unsigned m;
		key->x = (key->x + 1) & 255;
		key->y = (key->y + key->data[key->x]) & 255;
		xchg(key, key->x, key->y);
		m = key->data[(key->data[key->x] + key->data[key->y]) & 255];
		((unsigned char*)out)[i] = ((unsigned char*)in)[i] ^ m;
	}
}

/* encode a little endian sequence into a 32-bit value */
uint32_t rc4rand32(struct rc4_key *key)
{
	unsigned char buf[4] = { 0, 0, 0, 0 };
	rc4(key, sizeof buf, &buf, &buf);
	return buf[0] | ((uint32_t)buf[1] << 8)
		| ((uint32_t)buf[2] << 16) | ((uint32_t)buf[3] << 24);

}

/* encode a little endian sequence into a 64-bit value */
uint64_t rc4rand64(struct rc4_key *key)
{
	unsigned char buf[8] = { 0, 0, 0, 0, 0, 0, 0, 0 };
	rc4(key, sizeof buf, &buf, &buf);
	return buf[0] | ((uint64_t)buf[1] << 8)
		| ((uint64_t)buf[2] << 16) | ((uint64_t)buf[3] << 24)
		| ((uint64_t)buf[4] << 32) | ((uint64_t)buf[5] << 40)
		| ((uint64_t)buf[6] << 48) | ((uint64_t)buf[7] << 56);

}

static void hexdump(size_t len, const char *v)
{
	unsigned i;
	for (i = 0; i < len; i++) {
		printf("%02hhx", (unsigned)v[i]);
	}
	printf("\n");
}

int main()
{
	struct rc4_key key;
	char buf[] = "Attack at dawn";
	size_t buflen = strlen(buf);
	const char *secret = "Secret";

	/* exercise the crypto */
	rc4_set_key(&key, strlen(secret), secret);
	printf("key : %s\n", secret);
	printf("in  : %s\n", buf);
	rc4(&key, buflen, buf, buf);
	printf("out : ");
	hexdump(buflen, buf);

	/* use the engine as a random number generator */
	{
		int i;
		uint64_t m;
		for(i = 0; i < 20; i++) {
			m = rc4rand64(&key);
			printf("0x%016" PRIx64 "\n", m);
		}
	}
	return 0;
}