rc4.py

# -*- coding: utf-8 -*-
"""RC4

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1Vq1hsxIfe-PAS-DgSFxMSgbHivokf7XO
"""

class rc4:

  def __init__(self, plaintext, key):
    self.plaintext = plaintext
    self.key = key
    self.ciphertext = " "
    self.MOD = 256
    self.S = list(range(self.MOD))

  def KSA(self):
    ''' Key Scheduling Algorithm (from wikipedia):
        for i from 0 to 255
            S[i] := i
        endfor
        j := 0
        for i from 0 to 255
            j := (j + S[i] + key[i mod keylength]) mod 256
            swap values of S[i] and S[j]
        endfor
    '''
    key_length = len(self.key)
    # create the array "S"
    # [0,1,2, ... , 255]
    j = 0
    for i in range(self.MOD):
        j = (j + self.S[i] + self.key[i % key_length]) % self.MOD
        self.S[i], self.S[j] = self.S[j], self.S[i]  # swap values

    return self.S


  def PRGA(self):
    ''' Psudo Random Generation Algorithm (from wikipedia):
        i := 0
        j := 0
        while GeneratingOutput:
            i := (i + 1) mod 256
            j := (j + S[i]) mod 256
            swap values of S[i] and S[j]
            K := S[(S[i] + S[j]) mod 256]
            output K
        endwhile
    '''
    i = 0
    j = 0
    while True:
        i = (i + 1) % self.MOD
        j = (j + self.S[i]) % self.MOD

        self.S[i], self.S[j] = self.S[j], self.S[i]  # swap values
        K = self.S[(self.S[i] + self.S[j]) % self.MOD]
        yield K


  def get_keystream(self):
    ''' Takes the encryption key to get the keystream using PRGA
        return object is a generator
    '''
    self.S = self.KSA()
    return self.PRGA()
    

  def encrypt_logic(self):
    ''' :key -> encryption key used for encrypting, as hex string
        :text -> array of unicode values/ byte string to encrpyt/decrypt
    '''
    # For plaintext key, use this
    self.key = [ord(c) for c in self.key]
    # If key is in hex:
    # key = codecs.decode(key, 'hex_codec')
    # key = [c for c in key]
    keystream = self.get_keystream()

    res = []
    for c in self.plaintext:
        val = ("%02X" % (c ^ next(keystream)))  # XOR and taking hex
        res.append(val)
    return ''.join(res)
    

  def encrypt(self):
    ''' :key -> encryption key used for encrypting, as hex string
        :plaintext -> plaintext string to encrpyt
    '''
    self.plaintext = [ord(c) for c in self.plaintext]
    return self.encrypt_logic()

    
  def decrypt(self):
    ''' :key -> encryption key used for encrypting, as hex string
        :ciphertext -> hex encoded ciphered text using RC4
    '''
    self.ciphertext = codecs.decode(self.ciphertext, 'hex_codec')
    res = encrypt_logic(self.key, self.ciphertext)
    return codecs.decode(res, 'hex_codec').decode('utf-8')

cipher = rc4("Rachit", 'not-so-random-key')

cipher.encrypt()