RC4 implementation in Python

I started learning Python two months ago. To get the most out of the process, I decided to combine it with another interest of mine, cryptography, by trying to implement a very simple symmetric algorithm, RC4. Here is the code:

```#!/usr/bin/python3
#
# Author: Joao H de A Franco (jhafranco@acm.org)
#
# Description: RC4 implementation in Python 3
#
#          (CC BY-NC-SA 3.0)
#===========================================================

# Global variables
state = [None] * 256
p = q = None

def setKey(key):
"""RC4 Key Scheduling Algorithm (KSA)"""
global p, q, state
state = [n for n in range(256)]
p = q = j = 0
for i in range(256):
if len(key) > 0:
j = (j + state[i] + key[i % len(key)]) % 256
else:
j = (j + state[i]) % 256
state[i], state[j] = state[j], state[i]

def byteGenerator():
"""RC4 Pseudo-Random Generation Algorithm (PRGA)"""
global p, q, state
p = (p + 1) % 256
q = (q + state[p]) % 256
state[p], state[q] = state[q], state[p]
return state[(state[p] + state[q]) % 256]

def encrypt(inputString):
"""Encrypt input string returning a byte list"""
return [ord(p) ^ byteGenerator() for p in inputString]

def decrypt(inputByteList):
"""Decrypt input byte list returning a string"""
return "".join([chr(c ^ byteGenerator()) for c in inputByteList])```

To informally verify the correctness of this implementation, I wrote a separate Python module that uses Wikipedia’s RC4 test vectors for this purpose:

```import sys
import pyRC4

def main():
"""Verify the correctness of RC4 implementation using Wikipedia
test vectors"""

def intToList(inputNumber):
"""Convert a number into a byte list"""
inputString = "{:02x}".format(inputNumber)
return [int(inputString[i:i + 2], 16) for i in range(0, len(inputString), 2)]

def string_to_list(inputString):
"""Convert a string into a byte list"""
return [ord(c) for c in inputString]

def test(key, plaintext, ciphertext, testNumber):
success = True
pyRC4.setKey(string_to_list(key))
try:
assert pyRC4.encrypt(plaintext) == intToList(ciphertext)
print("RC4 encryption test #{:d} ok!".format(testNumber))
except AssertionError:
print("RC4 encryption test #{:d} failed".format(testNumber))
success = False
pyRC4.setKey(string_to_list(key))
try:
assert pyRC4.decrypt(intToList(ciphertext)) == plaintext
print("RC4 decryption test #{:d} ok!".format(testNumber))
except AssertionError:
print("RC4 decryption test #{:d} failed".format(testNumber))
success = False
return success

# Test vectors definition section
numberOfTests = 3
testVectorList = [{}] * numberOfTests

# Wikipedia test vector #1
testVectorList[0] = dict(key = "Key",
plaintext = "Plaintext",
testNumber = 1)

# Wikipedia test vector #2
testVectorList[1] = dict(key = "Wiki",
plaintext = "pedia",
ciphertext = 0x1021BF0420,
testNumber = 2)

# Wikipedia test vector #3
testVectorList[2] = dict(key = "Secret",
plaintext = "Attack at dawn",
ciphertext = 0x45A01F645FC35B383552544B9BF5,
testNumber = 3)

# Testing section
testSuccess = True
for p in range(numberOfTests):
testSuccess &= test(**testVectorList[p])
if testSuccess:
print("All RC4 tests succeeded!")
else:
print("At least one RC4 test failed")
sys.exit(1)
sys.exit()

if __name__ == '__main__':
main()
```