RC4 as iterator in Python

rc4.py

#!/usr/bin/env python2
# This is a simple RC4 python2 implementation as an iterator.
#
# Copyright (c) 2015 Hugo Guiroux: http://hugoguiroux.blogspot.fr/
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

# RC4 stream implemented as an iterator
class RC4:
    def __init__(self, key):
        # Init streams
        S = [i for i in xrange(256)]
        T = [ord(c) for c in (key * (256 / len(key)) + key[0:(256 % len(key))])]

        # First permutation of S
        j = 0
        for i in xrange(256):
            j = (j + S[i] + T[i]) % 256
            S[i], S[j] = S[j], S[i]

        # Internal state (no need to save T or the key)
        self.i, self.j = 0, 0
        self.S = S

    # Mandatory for Python iterator
    def __iter__(self):
        return self

    # The next method to act like an iterator
    def next(self):
        S, i, j = self.S, self.i, self.j

        # Change internal state by applying permutation on S
        i = (i + 1) % 256
        j = (j + S[i]) % 256
        S[i], S[j] = S[j], S[i]
        t = (S[i] + S[j]) % 256

        self.S, self.i, self.j = S, i, j

        return S[t]

# Simply XOR string with the RC4 stream (same function for decrypt & encrypt)
def crypt(S, string):
    return "".join([chr(ord(k) ^ next(S)) for k in string])

# Two different streams (one client side, one server side)
S = RC4("hidden_key")
D = RC4("hidden_key")

# E is the message encrypted sent to the server
E = crypt(S, "secret message")

# E is decrypted with D, the RC4 stream server side
print crypt(D, E)