#!/usr/bin/env python3
 
#   This code is public domain.
#
#   Improved ARC4 (IARC4) contains a number of proposed improvements over naive ARC4:
#
#   -   Uses KSA from VMPC minus an IV.
#   -   Uses 2 state spaces (RC4A). Splits the key and nonce to produce a key and nonce for each
#       state space. Each subkey and subnonce is XOR'd together to produce a new subkey.
#   -   Takes a nonce alongside the key. The key and nonce must be random and of even, equal
#       length, with 512 bytes per key/nonce suggested. **TODO**: They should be hashed, but they
#       are not currently, until I select a hash function with an appropriately sized output, which
#       won't limit the keyspace available to IARC4.
#   -   Drops the first 8192 (4096 per state space) iterations of the PRNG (RC4-drop8192).
#   -   A KeyExpiredError is raised after 255 iterations of the PRNG, excluding the initial drop.
#       Passing the `expires` option to IARC4 will alter this limit.
#
#   This code should not be considered secure. It has not been cryptanalyzed and should not be used
#   in production. This code is strictly experimental.
 
from bitstring import BitArray
from Crypto.Random import random
 
#   Generates a random n-byte string. Makes no checks to ensure uniqueness.
def random_bytes(n_bytes=512):
    return BitArray(uint=random.getrandbits(n_bytes*8), length=(n_bytes*8)).bytes
 
#   VMPC state preparation without vector
def prepare_state(key, state_size=256):
    key_size    = len(key)
    state       = [i for i in range(state_size)]
    j           = 0
 
    for n in range(state_size*3):
        i = n % state_size
        j = state[(((j + state[i]) % state_size) + key[i % key_size]) % state_size]
 
        state[i] ^= state[j]
        state[j] ^= state[i]
        state[i] ^= state[j]
 
    return state
 
#   Generic error (e.g., key and nonce of different lengths)
class KeyError(Exception): pass
 
#   Raised if a key/nonce combination has been used too many times already (255).
class KeyExpiredError(KeyError): pass
 
#   Uses 2 state spaces (RC4A)
def pseudorandom_generator(state_a, state_b, state_size=256, expires=8447):
    i       = 0
    j_a     = 0
    j_b     = 0
    n       = 0
 
    while True:
        n   += 1
 
        i   = (i + 1) % state_size
        j_a = (j_a + state_a[i]) % state_size
 
        state_a[i]   ^= state_a[j_a]
        state_a[j_a] ^= state_a[i]
        state_a[i]   ^= state_a[j_a]
 
        yield state_b[(state_a[i] + state_a[j_a]) % state_size]
 
        #   Limit the PRNG to 255 iterations after dropping 8192 iterations
        if (n == expires): break
        n += 1
 
        j_b = (j_b + state_b[i]) % state_size
 
        state_b[i]   ^= state_b[j_b]
        state_b[j_b] ^= state_b[i]
        state_b[i]   ^= state_b[j_b]
 
        yield state_a[(state_b[i] + state_b[j_b]) % state_size]
 
    raise KeyExpiredError
 
#   Main keystream generation and ciphering.
def IARC4(key, nonce, drop=8192, expires=255):
    key_length = len(key)
 
    #   Keys and nonces must be of even, equal lengths, 512 bytes apiece is recommended.
    if ((key_length != len(nonce)) or ((key_length % 2) != 0)):
        raise KeyError
 
    #   Split key and nonce in half to gain key_a, key_b, nonce_a, and nonce_b
    key_split   = key_length >> 1
    key_a       = BitArray(bytes=key[:key_split])
    key_b       = BitArray(bytes=key[key_split:])
    nonce_a     = BitArray(bytes=nonce[:key_split])
    nonce_b     = BitArray(bytes=nonce[key_split:])
 
    #   XOR keys and nonces.
    key_a = (key_a ^ nonce_a).bytes
    key_b = (key_b ^ nonce_b).bytes
 
    #   Prepare the PRNG
    state_a     = prepare_state(key_a)
    state_b     = prepare_state(key_b)
    keystream   = pseudorandom_generator(state_a, state_b, expires=(drop + expires))
 
    #   Discard first 4096 iterations of each state space
    for dropped in range(drop): next(keystream)
 
    #   Return the prepared PRNG to begin accepting input
    def _IARC4(text):
        #   Combine the keystream and the text stream
        for key_byte, text_byte in zip(keystream, text):
            yield key_byte ^ text_byte
 
    return _IARC4
 
#   Test Vector for IARC4
#
#   nonce:      b'\x01' + bytes(255) + b'\x02' + bytes(255)
#   key:        b'key_a' + bytes(251) + b'key_b' + bytes(251)
#   plaintext:  b'Plaintext'
#   ciphertext: b'\xb1\xc5\x8d)\x90\xf9WN\x94'
 
if (__name__ == "__main__"):
    nonce           = b'\x01' + bytes(255) + b'\x02' + bytes(255)
    key             = b'key_a' + bytes(251) + b'key_b' + bytes(251)
    plaintext       = b'Plaintext'
    stream_length   = len(plaintext)
    encrypter       = IARC4(key, nonce)
    encryption      = encrypter(plaintext)
    ciphertext      = bytes([next(encryption) for i in range(stream_length)])
    decrypter       = IARC4(key, nonce)
    decryption      = decrypter(ciphertext)
    plaintext       = bytes([next(decryption) for i in range(stream_length)])
 
    print("nonce:      %s\nkey:        %s\nplaintext:  %s\nciphertext: %s" %
    (nonce, key, plaintext, ciphertext))