1# Author: Trevor Perrin
2# See the LICENSE file for legal information regarding use of this file.
3
4"""Pure-Python RC4 implementation."""
5
6from .rc4 import RC4
7from .cryptomath import *
8
9def new(key):
10    return Python_RC4(key)
11
12class Python_RC4(RC4):
13    def __init__(self, keyBytes):
14        RC4.__init__(self, keyBytes, "python")
15        S = [i for i in range(256)]
16        j = 0
17        for i in range(256):
18            j = (j + S[i] + keyBytes[i % len(keyBytes)]) % 256
19            S[i], S[j] = S[j], S[i]
20
21        self.S = S
22        self.i = 0
23        self.j = 0
24
25    def encrypt(self, plaintextBytes):
26        ciphertextBytes = plaintextBytes[:]
27        S = self.S
28        i = self.i
29        j = self.j
30        for x in range(len(ciphertextBytes)):
31            i = (i + 1) % 256
32            j = (j + S[i]) % 256
33            S[i], S[j] = S[j], S[i]
34            t = (S[i] + S[j]) % 256
35            ciphertextBytes[x] ^= S[t]
36        self.i = i
37        self.j = j
38        return ciphertextBytes
39
40    def decrypt(self, ciphertext):
41        return self.encrypt(ciphertext)
42