RC4.sage

# key-scheduling algorithm

def KSA(key):
    
    keylength = len(key)
    S = range(256)
    j = 0
    
    for i in range(256):
        j = (j + S[i] + key[i % keylength]) % 256
        S[i], S[j] = S[j], S[i] # swap
    return S

# pseudo-random generation algorithm

def PRGA(S):
    
    i = 0
    j = 0
    while True:
        i = (i + 1) % 256
        j = (j + S[i]) % 256
        S[i], S[j] = S[j], S[i] # swap
        K = S[(S[i] + S[j]) % 256]
        yield K

def RC4(key):
    
    def convert_key(s):
        return [ord(c) for c in s]
    key = convert_key(key)
    
    S = KSA(key)
    return PRGA(S)