# -*- coding: utf-8 -*-
import string

class Trithemius(object):
    '''
    Classe de cifra Trithemius
    '''
    def __init__(self):
        self.plain = string.ascii_uppercase
        self.cipher = string.ascii_uppercase

    def shift(self):
        '''
        Desloca o alfabeto cifrado em uma casa.
        '''
        self.cipher = self.cipher[1:] + self.cipher[:1]

    def encrypt(self, text, decrypt=False):
        '''
        Cifra text com a cifra de Trithemius.
        Se decrypt é True, decifra ao invés de cifrar o text.
        '''
        text = text.replace(' ', '')
        ret_text = ''
        for char in text.upper():
            if decrypt:
                idx = self.cipher.find(char)
                ret_text += self.plain[idx]
            else:
                idx = self.plain.find(char)
                ret_text += self.cipher[idx]
            self.shift()
        self.cipher = string.ascii_uppercase
        return ret_text

    def decrypt(self, text):
        '''
        Decifra text com a cifra de Trithemius.
        '''
        return self.encrypt(text, True)