NOTE: I originally posted this on Snipplr.

This is a pure Python implementation of ARC4 as a generator to highlight its nature as a stream cipher. Several improvements can be made, for instance, it could take a nonce, use multiple state spaces (parallelizable), automatically discard the first 4K of the state space(s), use a more complex transformation than a simple swap, limit the # of bytes encrypted per nonce, etc.. The size of the state space is a parameter. The size of the key must not exceed the size of the state space, as additional key data will not be mixed into the prepared state.o

#!/usr/bin/env python3

def prepare_state(key, state_size=256):
    key_size    = len(key)
    state       = [i for i in range(state_size)]
    j           = 0

    for i in range(state_size):
        j = (j + state[i] + key[i % key_size]) % state_size

        state[i] ^= state[j]
        state[j] ^= state[i]
        state[i] ^= state[j]

    return state

def pseudorandom_generator(state):
    size    = len(state)
    i       = 0
    j       = 0

    while True:
        i = (i + 1) % size
        j = (j + state[i]) % size

        state[i] ^= state[j]
        state[j] ^= state[i]
        state[i] ^= state[j]

        yield state[(state[i] + state[j]) % size]

def ARC4(key, text):
    state       = prepare_state(key)
    keystream   = pseudorandom_generator(state)

    for key_byte, text_byte in zip(keystream, text):
        yield key_byte ^ text_byte

#   Test Vectors for ARC4
#   ---------------------
#   key:        b'Key'
#   plaintext:  b'Plaintext'
#   ciphertext: b'\xbb\xf3\x16\xe8\xd9@\xaf\n\xd3'
#   key:        b'Wiki'
#   plaintext:  b'pedia'
#   ciphertext: b'\x10!\xbf\x04 '
#   key:        b'Secret'
#   plaintext:  b'Attack at dawn'
#   ciphertext: b'E\xa0\x1fd_\xc3[85RTK\x9b\xf5'

#   Example usage
if (__name__ == "__main__"):
    from pprint import pprint as pp

    encryption  = ARC4(b'Key', b'Plaintext')
    ciphertext  = bytes([next(encryption) for i in range(len(b'Plaintext'))])

    decryption  = ARC4(b'Key', ciphertext)
    plaintext   = bytes([next(decryption) for i in range(len(ciphertext))])