from itertools import cycle
from easybits import Bits

class PolyCipher:
    """Implements a polyalphabetic cipher.
    The polyalphabetic cipher is like a Caesar cipher except that 
    the secret number changes for each character to be encrypted or
    decrypted. This makes frequency analysis much harder, because
    each plaintext space can be encrypted as a different character. 

    int_min and int_max represent the lowest and highest allowed int values
    of characters. They are set to include all the ASCII printable
    characters (https://en.wikipedia.org/wiki/ASCII#Printable_character_table)

    ASCII values outside this range (for example, '\n', the newline character),
    just get passed through unencrypted.
    """

    int_min = 32
    int_max = 127

    def __init__(self, secret):
        self.secret = secret

    def encrypt(self, plaintext):
        "Converts a plaintext message into an encrypted ciphertext"
        ciphertext = []
        for char, secret_char in zip(plaintext, cycle(self.secret)):
            plain_int = Bits(char, encoding='ascii').int
            if self.int_min <= plain_int and plain_int < self.int_max:
                secret_int = self.get_int(secret_char)
                cipher_int = self.rotate(secret_int, plain_int)
                ciphertext.append(Bits(cipher_int, length=8).ascii)
            else:
                ciphertext.append(char)
        return ''.join(ciphertext)

    def decrypt(self, ciphertext):
        "Converts an encrypted ciphertext into a plaintext message"
        plaintext = []
        for char, secret_char in zip(ciphertext, cycle(self.secret)):
            cipher_int = Bits(char, encoding='ascii').int
            if self.int_min <= cipher_int and cipher_int < self.int_max:
                secret_int = self.get_int(secret_char)
                plain_int = self.rotate(-secret_int, cipher_int)
                plaintext.append(Bits(plain_int, length=8).ascii)
            else:
                plaintext.append(char)
        return ''.join(plaintext)

    def rotate(self, secret, x):
        """Adds a secret number to x. 
        The modulo operator (%) is used to ensure that the result
        is greater than equal to int_min and less than int_max. 
        """
        range_size = self.int_max - self.int_min
        return (x + secret - self.int_min) % range_size + self.int_min

    def get_int(self, secret_char):
        """Converts an int or a single-character string into an int.
        When `secret_char` is an int, we just return it. Otherwise we
        return the character's ASCII value.
        """
        if isinstance(secret_char, int):
            return secret_char
        else:
            return Bits(secret_char, encoding='ascii').int