|
| 1 | +""" |
| 2 | +Wikipedia: https://en.wikipedia.org/wiki/Enigma_machine |
| 3 | +Video explanation: https://youtu.be/QwQVMqfoB2E |
| 4 | +Also check out Numberphile's and Computerphile's videos on this topic |
| 5 | +
|
| 6 | +This module contains function 'enigma' which emulates |
| 7 | +the famous Enigma machine from WWII. |
| 8 | +Module includes: |
| 9 | +- enigma function |
| 10 | +- showcase of function usage |
| 11 | +- 9 randnomly generated rotors |
| 12 | +- reflector (aka static rotor) |
| 13 | +- original alphabet |
| 14 | +
|
| 15 | +Created by TrapinchO |
| 16 | +""" |
| 17 | + |
| 18 | +# used alphabet -------------------------- |
| 19 | +# from string.ascii_uppercase |
| 20 | +abc = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ' |
| 21 | + |
| 22 | +# -------------------------- default selection -------------------------- |
| 23 | +# rotors -------------------------- |
| 24 | +rotor1 = 'EGZWVONAHDCLFQMSIPJBYUKXTR' |
| 25 | +rotor2 = 'FOBHMDKEXQNRAULPGSJVTYICZW' |
| 26 | +rotor3 = 'ZJXESIUQLHAVRMDOYGTNFWPBKC' |
| 27 | +# reflector -------------------------- |
| 28 | +reflector = {'A': 'N', 'N': 'A', 'B': 'O', 'O': 'B', 'C': 'P', 'P': 'C', 'D': 'Q', |
| 29 | + 'Q': 'D', 'E': 'R', 'R': 'E', 'F': 'S', 'S': 'F', 'G': 'T', 'T': 'G', |
| 30 | + 'H': 'U', 'U': 'H', 'I': 'V', 'V': 'I', 'J': 'W', 'W': 'J', 'K': 'X', |
| 31 | + 'X': 'K', 'L': 'Y', 'Y': 'L', 'M': 'Z', 'Z': 'M'} |
| 32 | + |
| 33 | +# -------------------------- extra rotors -------------------------- |
| 34 | +rotor4 = 'RMDJXFUWGISLHVTCQNKYPBEZOA' |
| 35 | +rotor5 = 'SGLCPQWZHKXAREONTFBVIYJUDM' |
| 36 | +rotor6 = 'HVSICLTYKQUBXDWAJZOMFGPREN' |
| 37 | +rotor7 = 'RZWQHFMVDBKICJLNTUXAGYPSOE' |
| 38 | +rotor8 = 'LFKIJODBEGAMQPXVUHYSTCZRWN' |
| 39 | +rotor9 = 'KOAEGVDHXPQZMLFTYWJNBRCIUS' |
| 40 | + |
| 41 | + |
| 42 | +def _validator(rotpos: tuple, rotsel: tuple, pb: str) -> tuple: |
| 43 | + """ |
| 44 | + Checks if the values can be used for the 'enigma' function |
| 45 | +
|
| 46 | + >>> _validator((1,1,1), (rotor1, rotor2, rotor3), 'POLAND') |
| 47 | + ((1, 1, 1), ('EGZWVONAHDCLFQMSIPJBYUKXTR', 'FOBHMDKEXQNRAULPGSJVTYICZW', \ |
| 48 | +'ZJXESIUQLHAVRMDOYGTNFWPBKC'), \ |
| 49 | +{'P': 'O', 'O': 'P', 'L': 'A', 'A': 'L', 'N': 'D', 'D': 'N'}) |
| 50 | +
|
| 51 | + :param rotpos: rotor_positon |
| 52 | + :param rotsel: rotor_selection |
| 53 | + :param pb: plugb -> validated and transformed |
| 54 | + :return: (rotpos, rotsel, pb) |
| 55 | + """ |
| 56 | + # Checks if there are 3 unique rotors |
| 57 | + |
| 58 | + unique_rotsel = len(set(rotsel)) |
| 59 | + if unique_rotsel < 3: |
| 60 | + raise Exception(f'Please use 3 unique rotors (not {unique_rotsel})') |
| 61 | + |
| 62 | + # Checks if rotor positions are valid |
| 63 | + rotorpos1, rotorpos2, rotorpos3 = rotpos |
| 64 | + if not 0 < rotorpos1 <= len(abc): |
| 65 | + raise ValueError(f'First rotor position is not within range of 1..26 (' |
| 66 | + f'{rotorpos1}') |
| 67 | + if not 0 < rotorpos2 <= len(abc): |
| 68 | + raise ValueError(f'Second rotor position is not within range of 1..26 (' |
| 69 | + f'{rotorpos2})') |
| 70 | + if not 0 < rotorpos3 <= len(abc): |
| 71 | + raise ValueError(f'Third rotor position is not within range of 1..26 (' |
| 72 | + f'{rotorpos3})') |
| 73 | + |
| 74 | + # Validates string and returns dict |
| 75 | + pb = _plugboard(pb) |
| 76 | + |
| 77 | + return rotpos, rotsel, pb |
| 78 | + |
| 79 | + |
| 80 | +def _plugboard(pbstring: str) -> dict: |
| 81 | + """ |
| 82 | + https://en.wikipedia.org/wiki/Enigma_machine#Plugboard |
| 83 | +
|
| 84 | + >>> _plugboard('PICTURES') |
| 85 | + {'P': 'I', 'I': 'P', 'C': 'T', 'T': 'C', 'U': 'R', 'R': 'U', 'E': 'S', 'S': 'E'} |
| 86 | + >>> _plugboard('POLAND') |
| 87 | + {'P': 'O', 'O': 'P', 'L': 'A', 'A': 'L', 'N': 'D', 'D': 'N'} |
| 88 | +
|
| 89 | + In the code, 'pb' stands for 'plugboard' |
| 90 | +
|
| 91 | + Pairs can be separated by spaces |
| 92 | + :param pbstring: string containing plugboard setting for the Enigma machine |
| 93 | + :return: dictionary containing converted pairs |
| 94 | + """ |
| 95 | + |
| 96 | + # tests the input string if it |
| 97 | + # a) is type string |
| 98 | + # b) has even length (so pairs can be made) |
| 99 | + if not isinstance(pbstring, str): |
| 100 | + raise TypeError(f'Plugboard setting isn\'t type string ({type(pbstring)})') |
| 101 | + elif len(pbstring) % 2 != 0: |
| 102 | + raise Exception(f'Odd number of symbols ({len(pbstring)})') |
| 103 | + elif pbstring == '': |
| 104 | + return {} |
| 105 | + |
| 106 | + pbstring.replace(' ', '') |
| 107 | + |
| 108 | + # Checks if all characters are unique |
| 109 | + tmppbl = set() |
| 110 | + for i in pbstring: |
| 111 | + if i not in abc: |
| 112 | + raise Exception(f'\'{i}\' not in list of symbols') |
| 113 | + elif i in tmppbl: |
| 114 | + raise Exception(f'Duplicate symbol ({i})') |
| 115 | + else: |
| 116 | + tmppbl.add(i) |
| 117 | + del tmppbl |
| 118 | + |
| 119 | + # Created the dictionary |
| 120 | + pb = {} |
| 121 | + for i in range(0, len(pbstring) - 1, 2): |
| 122 | + pb[pbstring[i]] = pbstring[i + 1] |
| 123 | + pb[pbstring[i + 1]] = pbstring[i] |
| 124 | + |
| 125 | + return pb |
| 126 | + |
| 127 | + |
| 128 | +def enigma(text: str, rotor_position: tuple, |
| 129 | + rotor_selection: tuple = (rotor1, rotor2, rotor3), plugb: str = '') -> str: |
| 130 | + """ |
| 131 | + The only difference with real-world enigma is that I allowed string input. |
| 132 | + All characters are converted to uppercase. (non-letter symbol are ignored) |
| 133 | + How it works: |
| 134 | + (for every letter in the message) |
| 135 | +
|
| 136 | + - Input letter goes into the plugboard. |
| 137 | + If it is connected to another one, switch it. |
| 138 | +
|
| 139 | + - Letter goes through 3 rotors. |
| 140 | + Each rotor can be represented as 2 sets of symbol, where one is shuffled. |
| 141 | + Each symbol from the first set has corresponding symbol in |
| 142 | + the second set and vice versa. |
| 143 | +
|
| 144 | + example: |
| 145 | + | ABCDEFGHIJKLMNOPQRSTUVWXYZ | e.g. F=D and D=F |
| 146 | + | VKLEPDBGRNWTFCJOHQAMUZYIXS | |
| 147 | +
|
| 148 | + - Symbol then goes through reflector (static rotor). |
| 149 | + There it is switched with paired symbol |
| 150 | + The reflector can be represented as2 sets, each with half of the alphanet. |
| 151 | + There are usually 10 pairs of letters. |
| 152 | +
|
| 153 | + Example: |
| 154 | + | ABCDEFGHIJKLM | e.g. E is paired to X |
| 155 | + | ZYXWVUTSRQPON | so when E goes in X goes out and vice versa |
| 156 | +
|
| 157 | + - Letter then goes through the rotors again |
| 158 | +
|
| 159 | + - If the letter is connected to plugboard, it is switched. |
| 160 | +
|
| 161 | + - Return the letter |
| 162 | +
|
| 163 | + >>> enigma('Hello World!', (1, 2, 1), plugb='pictures') |
| 164 | + 'KORYH JUHHI!' |
| 165 | + >>> enigma('KORYH, juhhi!', (1, 2, 1), plugb='pictures') |
| 166 | + 'HELLO, WORLD!' |
| 167 | + >>> enigma('hello world!', (1, 1, 1), plugb='pictures') |
| 168 | + 'FPNCZ QWOBU!' |
| 169 | + >>> enigma('FPNCZ QWOBU', (1, 1, 1), plugb='pictures') |
| 170 | + 'HELLO WORLD' |
| 171 | +
|
| 172 | +
|
| 173 | + :param text: input message |
| 174 | + :param rotor_position: tuple with 3 values in range 1..26 |
| 175 | + :param rotor_selection: tuple with 3 rotors () |
| 176 | + :param plugb: string containing plugboard configuration (default '') |
| 177 | + :return: en/decrypted string |
| 178 | + """ |
| 179 | + |
| 180 | + text = text.upper() |
| 181 | + rotor_position, rotor_selection, plugboard = _validator( |
| 182 | + rotor_position, rotor_selection, plugb.upper()) |
| 183 | + |
| 184 | + rotorpos1, rotorpos2, rotorpos3 = rotor_position |
| 185 | + rotor1, rotor2, rotor3 = rotor_selection |
| 186 | + rotorpos1 -= 1 |
| 187 | + rotorpos2 -= 1 |
| 188 | + rotorpos3 -= 1 |
| 189 | + plugboard = plugboard |
| 190 | + |
| 191 | + result = [] |
| 192 | + |
| 193 | + # encryption/decryption process -------------------------- |
| 194 | + for symbol in text: |
| 195 | + if symbol in abc: |
| 196 | + |
| 197 | + # 1st plugboard -------------------------- |
| 198 | + if symbol in plugboard: |
| 199 | + symbol = plugboard[symbol] |
| 200 | + |
| 201 | + # rotor ra -------------------------- |
| 202 | + index = abc.index(symbol) + rotorpos1 |
| 203 | + symbol = rotor1[index % len(abc)] |
| 204 | + |
| 205 | + # rotor rb -------------------------- |
| 206 | + index = abc.index(symbol) + rotorpos2 |
| 207 | + symbol = rotor2[index % len(abc)] |
| 208 | + |
| 209 | + # rotor rc -------------------------- |
| 210 | + index = abc.index(symbol) + rotorpos3 |
| 211 | + symbol = rotor3[index % len(abc)] |
| 212 | + |
| 213 | + # reflector -------------------------- |
| 214 | + # this is the reason you don't need another machine to decipher |
| 215 | + |
| 216 | + symbol = reflector[symbol] |
| 217 | + |
| 218 | + # 2nd rotors |
| 219 | + symbol = abc[rotor3.index(symbol) - rotorpos3] |
| 220 | + symbol = abc[rotor2.index(symbol) - rotorpos2] |
| 221 | + symbol = abc[rotor1.index(symbol) - rotorpos1] |
| 222 | + |
| 223 | + # 2nd plugboard |
| 224 | + if symbol in plugboard: |
| 225 | + symbol = plugboard[symbol] |
| 226 | + |
| 227 | + # moves/resets rotor positions |
| 228 | + rotorpos1 += 1 |
| 229 | + if rotorpos1 >= len(abc): |
| 230 | + rotorpos1 = 0 |
| 231 | + rotorpos2 += 1 |
| 232 | + if rotorpos2 >= len(abc): |
| 233 | + rotorpos2 = 0 |
| 234 | + rotorpos3 += 1 |
| 235 | + if rotorpos3 >= len(abc): |
| 236 | + rotorpos3 = 0 |
| 237 | + |
| 238 | + # else: |
| 239 | + # pass |
| 240 | + # Error could be also raised |
| 241 | + # raise ValueError( |
| 242 | + # 'Invalid symbol('+repr(symbol)+')') |
| 243 | + result.append(symbol) |
| 244 | + |
| 245 | + return "".join(result) |
| 246 | + |
| 247 | + |
| 248 | +if __name__ == '__main__': |
| 249 | + message = 'This is my Python script that emulates the Enigma machine from WWII.' |
| 250 | + rotor_pos = (1, 1, 1) |
| 251 | + pb = 'pictures' |
| 252 | + rotor_sel = (rotor2, rotor4, rotor8) |
| 253 | + en = enigma(message, rotor_pos, rotor_sel, pb) |
| 254 | + |
| 255 | + print('Encrypted message:', en) |
| 256 | + print('Decrypted message:', enigma(en, rotor_pos, rotor_sel, pb)) |
0 commit comments