#!/usr/bin/python3
# -*- coding: utf-8 -*-

from itertools import cycle
from typing import List, Generator


# Key preparation:
def key_schedule(key: str) -> List[int]:
    state_array = list(range(256))
    key_bytes = cycle(ord(x) for x in key)
    j = 0
    for i in range(256):
        j = (j + state_array[i] + next(key_bytes)) % 256
        state_array[i], state_array[j] = state_array[j], state_array[i]
    return state_array


# Keystream generator:
def pseudorandom_generator(state_array: List[int]) -> Generator[int, None, None]:
    j = 0
    for i in cycle(range(256)):
        j = (j + state_array[i]) % 256
        state_array[i], state_array[j] = state_array[j], state_array[i]
        k = (state_array[i] + state_array[j]) % 256
        yield state_array[k]


def main() -> None:
    plaintext = "secret"
    print(f"Plaintext is: {plaintext}\n")

    # encrypt
    print("Encrypting now, one byte at a time:")
    key = "not-so-random-key"
    state_array = key_schedule(key)
    infinite_key_generator = pseudorandom_generator(state_array)
    translate_array = []
    for c in plaintext:
        print("%02X" % (ord(c) ^ next(infinite_key_generator)))
        translate_array.append("%02X" % (ord(c) ^ next(infinite_key_generator)))
    ciphertext = "".join([chr(int(c, 16)) for c in translate_array])
    print(f"\nCiphertext is: {ciphertext}\n")

    # decrypt
    print("Decrypting now, one byte at a time:")
    key = "not-so-random-key"
    state_array = key_schedule(key)
    infinite_key_generator = pseudorandom_generator(state_array)
    translate_array = []
    for c in ciphertext:
        print("%02X" % (ord(c) ^ next(infinite_key_generator)))
        translate_array.append("%02X" % (ord(c) ^ next(infinite_key_generator)))
    plaintext = "".join(chr(int(c, 16)) for c in translate_array)
    print(f"\nPlaintext is: {plaintext}\n")


if __name__ == "__main__":
    main()