Create a gist now

Instantly share code, notes, and snippets.

What would you like to do?
RC4 encryption
from hashlib import sha1
from itertools import islice, imap
from operator import xor, eq
DROP_N = 1536
# drop 1536 bytes is adviced in RFC 4345
# http://tools.ietf.org/html/rfc4345
def _drop(it, n):
try:
islice(it, n, n).next()
except StopIteration:
pass
return it
def _key_sched(key):
S = range(256)
j = 0
key = map(ord, key)
for i in xrange(256):
j = (j + S[i] + key[i % len(key)]) % 256
S[i], S[j] = S[j], S[i]
return S
def _key_stream(key):
S = _key_sched(key)
i = 0
j = 0
while True:
i = (i + 1) % 256
j = (j + S[i]) % 256
S[i], S[j] = S[j], S[i]
t = (S[i] + S[j]) % 256
yield S[t]
class RC4(object):
def __init__(self, shared_key, iv, drop_n=DROP_N):
key = sha1(iv + shared_key).digest()
self.stream = _drop(_key_stream(key), drop_n)
def __call__(self, data):
bytes = imap(ord, data)
converted_bytes = imap(xor, bytes, self.stream)
return ''.join(imap(chr, converted_bytes))
def test_official_vectors():
# from http://tools.ietf.org/html/rfc6229
v = {'0102030405060708':
[(0, '97ab8a1bf0afb96132f2f67258da15a8'),
(1536, '8369e1a965610be887fbd0c79162aafb'),
(4096, '3fb46e2c8f0f068edce8cdcd7dfc5862')],
'0102030405060708090a0b0c0d0e0f101112131415161718191a1b1c1d1e1f20':
[(0, 'eaa6bd25880bf93d3f5d1e4ca2611d91'),
(1536, '3e34135c79db010200767651cf263073'),
(4096, 'f3e4c0a2e02d1d01f7f0a74618af2b48')],
'ebb46227c6cc8b37641910833222772a':
[(0, '720c94b63edf44e131d950ca211a5a30'),
(1536, '1175da6ee756de46a53e2b075660b770'),
(4096, '5bbeb4787d59e5373fdbea6c6f75c29b')],
'c109163908ebe51debb46227c6cc8b37641910833222772a':
[(0, '54b64e6b5a20b5e2ec84593dc7989da7'),
(1536, 'baefe6d9ceb651842260e0d1e05e3b90'),
(4096, 'e4dd2e98d6960fae0b43545456743391')]
}
try:
all = all
except NameError:
def all(iterable):
for element in iterable:
if not element:
return False
return True
for key, data in v.iteritems():
key = key.decode('hex')
for drop, expected in data:
stream = _drop(_key_stream(key), drop)
expected = [ord(c) for c in expected.decode('hex')]
assert all(imap(eq, stream, expected))
def test_simple():
plain = 'ABCD123'
cipher = RC4('shared_secret', 'iv')(plain)
assert plain == RC4('shared_secret', 'iv')(cipher)
if __name__ == '__main__':
test_official_vectors()
test_simple()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment