cypher.py

import random
import copy

alphabet = list('ABCDEFGHIJKLMNOPQRSTUVWXYZ')
scrambled_alphabet = list('ABCDEFGHIJKLMNOPQRSTUVWXYZ')
random.seed(0)
random.shuffle(scrambled_alphabet)
original_rotor1 = [(s, a) for a, s in zip(alphabet, scrambled_alphabet)]
random.shuffle(scrambled_alphabet)
original_rotor2 = [(s, a) for a, s in zip(alphabet, scrambled_alphabet)]
random.shuffle(scrambled_alphabet)
original_rotor3 = [(s, a) for a, s in zip(alphabet, scrambled_alphabet)]


def rotate(rotor, n):
    scrambled_alphabet = [x[0] for x in rotor]
    rotated_scrambled_alphabet = scrambled_alphabet[n:] + scrambled_alphabet[:n]
    return [(s, a) for a, s in zip(alphabet, rotated_scrambled_alphabet)]

def enigma(message, key):
    key_inds = [alphabet.index(k) for k in key]
    rotor1 = copy.deepcopy(original_rotor1)
    rotor2 = copy.deepcopy(original_rotor2)
    rotor3 = copy.deepcopy(original_rotor3)
    rotor1 = rotate(rotor1, key_inds[0])
    rotor2 = rotate(rotor2, key_inds[1])
    rotor3 = rotate(rotor3, key_inds[2])
    position = 0
    count_rotor_1 = 0
    count_rotor_2 = 0
    count_rotor_3 = 0
    encrypted_message = ''
    while position < len(message):
        inv_rotor1 = sorted([(a, s) for a, s in zip(alphabet, rotor1)], key=lambda x: x[1])
        inv_rotor2 = sorted([(a, s) for a, s in zip(alphabet, rotor2)], key=lambda x: x[1])
        inv_rotor3 = sorted([(a, s) for a, s in zip(alphabet, rotor3)], key=lambda x: x[1])
        character = alphabet.index(message[position])
        character = alphabet.index(rotor1[character][0])
        character = alphabet.index(rotor2[character][0])
        character = alphabet.index(rotor3[character][0])
        character = 25 - character
        character = alphabet.index(inv_rotor3[character][0])
        character = alphabet.index(inv_rotor2[character][0])
        character = alphabet.index(inv_rotor1[character][0])
        encrypted_message += alphabet[character]
        rotor1 = rotate(rotor1, 1)
        if position % 26 == 25:
            rotor2 = rotate(rotor2, 1)
            if position % (26 * 26) == ((26 * 26) - 1):
                rotor3 = rotate(rotor3, 1)
        position += 1
    return encrypted_message


def caesar(message, key):
    encrypted_message = []
    for letter in message:
        letter_ind = alphabet.index(letter)
        new_letter_ind = (letter_ind + key) % 26
        new_letter = alphabet[new_letter_ind]
        encrypted_message.append(new_letter)
    return ''.join(encrypted_message)