diff --git a/src/vigenere_cipher_helper.py b/src/vigenere_cipher_helper.py index fb5401a..b560a7f 100644 --- a/src/vigenere_cipher_helper.py +++ b/src/vigenere_cipher_helper.py @@ -17,11 +17,16 @@ class VigenereCipher(object): new_text = [] for l_text, l_key in zip(text, new_key): - if l_key in self.alphabet: - new_pos = self.alphabet.index(l_text) + self.alphabet.index(l_key) + if l_text in self.alphabet: + if l_key in self.alphabet: + new_pos = ( + self.alphabet.index(l_text) + self.alphabet.index(l_key) + ) % len(self.alphabet) + else: + new_pos = self.alphabet.index(l_text) + new_text.append(self.alphabet[new_pos]) else: - new_pos = self.alphabet.index(l_text) - new_text.append(self.alphabet[new_pos]) + new_text.append(l_text) return "".join(new_text) def decode(self, text): @@ -37,18 +42,15 @@ class VigenereCipher(object): new_text = [] - for l_text in text: - for i, l_abc in enumerate(self.alphabet): - if l_text == l_abc: - for l_key in new_key: - for j, l_abc in enumerate(self.alphabet): - if l_key == l_abc: - new_pos = (i - j) % len(self.alphabet) - else: - new_pos = i + for l_text, l_key in zip(text, new_key): + if l_text in self.alphabet: + if l_key in self.alphabet: + new_pos = ( + self.alphabet.index(l_text) - self.alphabet.index(l_key) + ) % len(self.alphabet) else: - continue - for i, l_abc in enumerate(self.alphabet): - if new_pos == i: - new_text.append(l_abc) + new_pos = self.alphabet.index(l_text) + new_text.append(self.alphabet[new_pos]) + else: + new_text.append(l_text) return "".join(new_text)