KnowHow

import string


class VigenereCipher(object):
    def __init__(self, message: str, key: str) -> None:
        self.message = message
        self._len_message = len(message)
        self.key = self._generate_keys(key)
        self._len_alphabet = ord('Z') - ord('A') + 1

    def _generate_keys(self, keyword: str) -> str:
        key_string = ''
        for i in range(self._len_message):
            key_string += keyword[i % len(keys)]
        return key_string

    def fetch_keys_index(self, index) -> int:
        if self.key[index].isupper():
            key_index = ord(self.key[index]) - ord('A')
        else:
            key_index = ord(self.key[index]) - ord('a')
        return key_index

    def encrypt(self) -> str:
        result = ''
        for idx, char in enumerate(self.message):
            if char.isupper():
                char_index = ord(char) - ord('A')
                key_index = self.fetch_keys_index(idx)
                encrypt_index = ((char_index + key_index) % self._len_alphabet) + ord('A')
            elif char.islower():
                char_index = ord(char) - ord('a')
                key_index = self.fetch_keys_index(idx)
                encrypt_index = ((char_index + key_index) % self._len_alphabet) + ord('a')
            else:
                encrypt_index = ord(char)
            result += chr(encrypt_index)
        return result

    def decrypt(self) -> str:
        result = ''
        for idx, char in enumerate(self.message):
            if char.isupper():
                char_index = ord(char) - ord('A')
                key_index = self.fetch_keys_index(idx)
                encrypt_index = ((char_index - key_index + self._len_alphabet) % self._len_alphabet) + ord('A')
            elif char.islower():
                char_index = ord(char) - ord('a')
                key_index = self.fetch_keys_index(idx)
                encrypt_index = ((char_index - key_index + self._len_alphabet) % self._len_alphabet) + ord('a')
            else:
                encrypt_index = ord(char)
            result += chr(encrypt_index)
        return result


if __name__ == '__main__':
    # s = "AaBbCcDdZz"
    s = "ABCDZ"
    keys = 'HELLO'
    print(s)
    print(encode_vigenere_cipher(s, keys))
    # print(generate_key('ABCDEFGZ', 'CAT'))
    print(encrypt(s, keys))

    print('object test 1')
    vg_en = VigenereCipher(s, keys)
    ee = vg_en.encrypt()
    print(ee)
    vg_d = VigenereCipher(ee, keys)
    dd = vg_d.decrypt()
    print(dd)
    print('+'*10)

    t = 'ATTACK SILICON VALLEY'
    k = generate_key(t, 'HELLO')
    e = encrypt(t, k)
    print(e)
    print(encode_vigenere_cipher(t, 'HELLO'))

    d = decrypt(e, k)
    print(d)

    print('object test 2')
    print(t)
    key = 'HELLO'
    vg_e2 = VigenereCipher(t, key)
    e2 = vg_e2.encrypt()
    print(e2)
    vg_d2 = VigenereCipher(e2, key)
    d2 = vg_d2.decrypt()
    print(d2)

    print('object test 3')
    t = 'Attack Silicon Valley'
    print(t)
    key = 'HELLO'
    vg_e2 = VigenereCipher(t, key)
    e2 = vg_e2.encrypt()
    print(e2)
    vg_d2 = VigenereCipher(e2, key)
    d2 = vg_d2.decrypt()
    print(d2)