rc4.py

'''
rc4.py

RC4 is a variable-key-size stream cipher developed in 1987 by Ron Rivest for 
RSA Data Security, Inc.
'''

__all__ = ['crypt']

def sbox(key):
	"""Initialize the S-box.
	>>> len(sbox("asdfg"))
	256
	"""
	# 1st fill it linearly.
	s = range(0, 256)
	j = 0
	for i in range(0, 256):
		j = (j + s[i] + ord(key[i % len(key)])) % 256
		s[i], s[j] = s[j], s[i]
	return s

def crypt(data, key):
	"""Encrypt/Decrypt using RC4.
	>>> t = "hello kitty"
	>>> p = "less"
	>>> s = crypt(t, p)
	>>> crypt(s, p) == t
	True
	"""
	s = sbox(key)
	i = j = 0
	out = []
	for ch in data:
		i = (i + 1) % 256
		j = (j + s[i]) % 256
		s[i], s[j] = s[j], s[i]
		t = (s[i] + s[j]) % 256
		out.append(chr(ord(ch) ^ s[t]))
	return ''.join(out)

if __name__ == "__main__":
	import doctest
	doctest.testmod()