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
#
# License: Attribution-NonCommercial-ShareAlike 3.0 Unported
#          (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",
                             ciphertext = 0xBBF316E8D940AF0AD3,
                             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()

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out /  Change )

Facebook photo

You are commenting using your Facebook account. Log Out /  Change )

Connecting to %s