import random, base64

__all__ = ['crypt', 'encrypt', 'decrypt']

def crypt(data, key):
    """RC4 algorithm"""
    x = 0
    box = range(256)
    for i in range(256):
        x = (x + box[i] + ord(key[i % len(key)])) % 256
        box[i], box[x] = box[x], box[i]
    x = y = 0
    out = []
    for char in data:
        x = (x + 1) % 256
        y = (y + box[x]) % 256
        box[x], box[y] = box[y], box[x]
        out.append(chr(ord(char) ^ box[(box[x] + box[y]) % 256]))
    
    return ''.join(out)

def encrypt(data, key, encode=base64.standard_b64encode, salt_length=8):
    """RC4 encryption with random salt and final encoding"""
    salt = ''
    for n in range(salt_length):
        salt += chr(random.randrange(256))
    data = salt + crypt(data, key + salt)
    if encode:
        data = encode(data)
    return data
    
def decrypt(data, key, decode=base64.standard_b64decode, salt_length=8):
    """RC4 decryption of encoded data"""
    if decode:
        data = decode(data)
    salt = data[:salt_length]
    return crypt(data[salt_length:], key + salt)

if __name__ == '__main__':
    for i in range(10):
        data = encrypt('secret message', 'my-key')
        print data
        print decrypt(data, 'my-key')