module std::crypto::rc4;
// Copyright (c) 2021 Christoffer Lerno. All rights reserved.
// Use of this source code is governed by the MIT license
// a copy of which can be found in the LICENSE_STDLIB file.

struct Rc4
{
	uint i, j;
	char[256] state;
}

/**
 * Initialize the RC4 state.
 *
 * @param [in] key "The key to use"
 * @require key.len > 0 "The key must be at least 1 byte long"
 **/
fn void Rc4.init(&self, char[] key)
{
	// Init the state matrix
	foreach (char i, &c : self.state) *c = i;
	for (int i = 0, int j = 0; i < 256; i++)
	{
		j = (j + self.state[i] + key[i % key.len]) & 0xFF;
		@swap(self.state[i], self.state[j]);
	}
	self.i = 0;
	self.j = 0;
}

/**
 * Run a single pass of en/decryption using a particular key.
 * @param [in] key
 * @param [inout] data
 **/
fn void crypt(char[] key, char[] data)
{
	Rc4 rc4;
	rc4.init(key);
	rc4.crypt(data, data);
}

/**
 * Encrypt or decrypt a sequence of bytes.
 *
 * @param [in] in "The input"
 * @param [out] out "The output"
 * @require in.len <= out.len "Output would overflow"
 **/
fn void Rc4.crypt(&self, char[] in, char[] out)
{
	uint i = self.i;
	uint j = self.j;
	char* state = &self.state;
	isz len = in.len;
	foreach (idx, c : in)
	{
		i = (i + 1) & 0xFF;
		j = (j + state[i]) & 0xFF;
		@swap(state[i], state[j]);
		out[idx] = in[idx] ^ state[(state[i] + state[j]) & 0xFF];
	}
	self.i = i;
	self.j = j;
}

/**
 * Clear the rc4 state.
 *
 * @param [&out] self "The RC4 State"
 **/
fn void Rc4.destroy(&self)
{
	*self = {};
}