RC4 encryption

rc4.py

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()