Simple RC4 Stream Cipher implementation in Lua

rc4.lua

-- small Lua RC4 library.
-- Tested in luajit and lua5.3.
-- Use for testing purposes only!
-- RC4 is broken, and this is not a well-tested implementation.
local ok,bit = pcall(require, "bit")
if not ok then
	bit = require("bit32")
end
assert(bit and (bit.bxor(0xFF,0xF0)==0x0F)) -- test bit library xor

local rc4 = {}

-- Swap two table values based on their index.
local function swap(t,ia,ib)
	local a = t[ia]
	local b = t[ib]
	t[ia] = b
	t[ib] = a
end

-- Dump a state(S) in human-readable form as a string(also valid lua table)
function rc4.dump_S(S)
	local str = {}
	for i=0, 255 do
		str[i+1] = ("%.2X"):format(S[i])
	end
	return ("{\n\ti = %d,\n\tj = %d,\n\tdata = %q\n}"):format(S.i, S.j, table.concat(str))
end

-- Key-scheduling algorithm.
-- Based on the key, generate the initial state(S)
-- Returns new RC4 state(S)
function rc4.ksa(key)
	local S = { i = 0, j = 0}
	for i=0, 255 do
		S[i] = i
	end
	local j = 0
	local key_len = #key
	for i=0, 255 do
		j = (j + S[i] + string.byte(key, (i%key_len)+1)) % 256
		swap(S, i,j)
	end
	return S
end

-- Pseudo-random number generator.
-- Based on the state(S) return the next pseudo-random number(0-255), and iterate the state.
-- This sequence of numbers is used as the keystream for encryption/decryption.
function rc4.prng(S)
	S.i = (S.i+1) % 256
	S.j = (S.j+S[S.i]) % 256
	S.i, S.j = S.j, S.i
	swap(S, "i","j")
	return S[(S[S.i]+S[S.j]) % 256]
end

-- Generate the initial keystate using the ksa,
-- then iterate over every byte in the plaintext,
-- and xor it with the keystream byte returned by the prng iteration.
function rc4:encrypt(key, plaintext)
	assert(type(key)=="string")
	assert(type(plaintext)=="string")
	assert(#key>0)
	assert(#key<=16)
	local cyphertext = {}
	local S = self.ksa(key) -- generate the initial state based on the key
	for i=1, #plaintext do
		local prn = self.prng(S) -- get a random number by iterating the prng state
		cyphertext[i] = string.char(bit.bxor(prn, string.byte(plaintext, i)))
	end
	return table.concat(cyphertext)
end

-- Decrypt works the same as encrypt:
-- plaintext = rc4:decrypt(key, cyphertext)
rc4.decrypt = rc4.encrypt

return rc4