From 0bd56b3405d997f00acf2c484bebbc655896edc0 Mon Sep 17 00:00:00 2001 From: DingJunyao Date: Tue, 7 Dec 2021 00:50:45 +0800 Subject: [PATCH] v1.0, raw --- .idea/.gitignore | 3 + .idea/enigma.iml | 10 ++ .../inspectionProfiles/profiles_settings.xml | 6 + .idea/misc.xml | 4 + .idea/modules.xml | 8 + .idea/vcs.xml | 6 + README.md | 122 ++++++++++++++++ pyproject.toml | 6 + setup.cfg | 28 ++++ src/__init__.py | 0 src/enigma/__init__.py | 1 + src/enigma/enigma.py | 100 +++++++++++++ src/enigma/part/__init__.py | 1 + src/enigma/part/plate/__init__.py | 3 + src/enigma/part/plate/entry_plate.py | 13 ++ src/enigma/part/plate/plate.py | 138 ++++++++++++++++++ src/enigma/part/plate/reflector.py | 23 +++ src/enigma/part/plate/rotor.py | 85 +++++++++++ src/enigma/part/plugboard.py | 29 ++++ src/enigma/sample_plate/__init__.py | 4 + src/enigma/sample_plate/reflector.py | 6 + src/enigma/sample_plate/rotor.py | 12 ++ tests/test_enigma.py | 70 +++++++++ tests/test_enigma_plug.py | 50 +++++++ tests/test_plate.py | 6 + 25 files changed, 734 insertions(+) create mode 100644 .idea/.gitignore create mode 100644 .idea/enigma.iml create mode 100644 .idea/inspectionProfiles/profiles_settings.xml create mode 100644 .idea/misc.xml create mode 100644 .idea/modules.xml create mode 100644 .idea/vcs.xml create mode 100644 pyproject.toml create mode 100644 setup.cfg create mode 100644 src/__init__.py create mode 100644 src/enigma/__init__.py create mode 100644 src/enigma/enigma.py create mode 100644 src/enigma/part/__init__.py create mode 100644 src/enigma/part/plate/__init__.py create mode 100644 src/enigma/part/plate/entry_plate.py create mode 100644 src/enigma/part/plate/plate.py create mode 100644 src/enigma/part/plate/reflector.py create mode 100644 src/enigma/part/plate/rotor.py create mode 100644 src/enigma/part/plugboard.py create mode 100644 src/enigma/sample_plate/__init__.py create mode 100644 src/enigma/sample_plate/reflector.py create mode 100644 src/enigma/sample_plate/rotor.py create mode 100644 tests/test_enigma.py create mode 100644 tests/test_enigma_plug.py create mode 100644 tests/test_plate.py diff --git a/.idea/.gitignore b/.idea/.gitignore new file mode 100644 index 0000000..26d3352 --- /dev/null +++ b/.idea/.gitignore @@ -0,0 +1,3 @@ +# Default ignored files +/shelf/ +/workspace.xml diff --git a/.idea/enigma.iml b/.idea/enigma.iml new file mode 100644 index 0000000..c6d5d10 --- /dev/null +++ b/.idea/enigma.iml @@ -0,0 +1,10 @@ + + + + + + + + + + \ No newline at end of file diff --git a/.idea/inspectionProfiles/profiles_settings.xml b/.idea/inspectionProfiles/profiles_settings.xml new file mode 100644 index 0000000..105ce2d --- /dev/null +++ b/.idea/inspectionProfiles/profiles_settings.xml @@ -0,0 +1,6 @@ + + + + \ No newline at end of file diff --git a/.idea/misc.xml b/.idea/misc.xml new file mode 100644 index 0000000..80bf3eb --- /dev/null +++ b/.idea/misc.xml @@ -0,0 +1,4 @@ + + + + \ No newline at end of file diff --git a/.idea/modules.xml b/.idea/modules.xml new file mode 100644 index 0000000..3e579c9 --- /dev/null +++ b/.idea/modules.xml @@ -0,0 +1,8 @@ + + + + + + + + \ No newline at end of file diff --git a/.idea/vcs.xml b/.idea/vcs.xml new file mode 100644 index 0000000..94a25f7 --- /dev/null +++ b/.idea/vcs.xml @@ -0,0 +1,6 @@ + + + + + + \ No newline at end of file diff --git a/README.md b/README.md index 779bbe9..74b6b06 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,124 @@ # enigma Python-based Enigma. + +## Install + +```bash +pip install enigma +``` + +## Usage + +### Import + +```python +from enigma import Enigma +``` + +The package also contains some sample plate: + +```python +from enigma.sample_plate.rotor import rotor_I, rotor_II, rotor_III +from enigma.sample_plate.reflector import reflector_B +``` +### Defining + +Make sure the rotors are from right to left: + +```python +e = Enigma([rotor_III(), rotor_II(), rotor_I()], reflector_B()) +``` + +### Encryption + +Just input: + +```python +# Encryption +assert e.input('HELLOWORLDBYTHEAUTHOROFTHISPACKAGEDINGJUNYAO') == 'ILBDAAMTAZIJNXSCSIJJPDDWZBRCCUPGQXGRJXQOFGHL' +``` + +And change the positions of the rotors (Defaults to all origin). + +```python +e.set_position() +assert e.input('ILBDAAMTAZIJNXSCSIJJPDDWZBRCCUPGQXGRJXQOFGHL') == 'HELLOWORLDBYTHEAUTHOROFTHISPACKAGEDINGJUNYAO' +``` + +Number starts from 0, and also from right to left. You can also use letter. + +```python +e.set_position(3, 12, 21) +assert e.input('HELLOWORLDBYTHEAUTHOROFTHISPACKAGEDINGJUNYAO') == 'XTGHAGDIVUPGBZVQSFMBSGLKVQHQWESYRTSRMOOFGRLE' +e.set_position('D', 12, 21) +assert e.input('XTGHAGDIVUPGBZVQSFMBSGLKVQHQWESYRTSRMOOFGRLE') == 'HELLOWORLDBYTHEAUTHOROFTHISPACKAGEDINGJUNYAO' +``` + +### Customize + +You can freely customize your Enigma. For example, customize the circuits of rotors: + +```python +from enigma.part.plate import Rotor, Reflector + +rotor_I = Rotor('EKMFLGDQVZNTOWYHXUSPAIBRCJ', name='Rotor I', turnover='Q') +rotor_II = Rotor('AJDKSIRUXBLHWTMCQGZNPYFVOE', name='Rotor II', turnover='E') +rotor_III = Rotor('BDFHJLCPRTXVZNYEIWGAKMUSQO', name='Rotor III', turnover='V') +reflector_B = Reflector('YRUHQSLDPXNGOKMIEBFZCWVJAT', name='Reflector B') +e = Enigma([rotor_III(), rotor_II(), rotor_I()], reflector_B()) +# same as e above +``` + +A tiny Enigma: + +```python +map_source = 'ASDF' +e_custom_1 = Enigma( + [ + Rotor('AFSD', name='I', map_source=map_source), + Rotor('SDAF', name='II', map_source=map_source), + Rotor('DFAS', name='III', map_source=map_source), + ], + Reflector('DFAS', name='R', map_source=map_source), + rotate_up=True, rotate_after_type=True +) +assert e_custom_1.input('AA') == 'SD' +e_custom_1.set_position() +assert e_custom_1.input('SD') == 'AA' +``` + +And even in other character! + +```python +map_source_2 = '甲乙丙丁' +e_custom_2 = Enigma( + [ + Rotor('甲丁乙丙', name='I', map_source=map_source_2), + Rotor('乙丙甲丁', name='II', map_source=map_source_2), + Rotor('丙丁甲乙', name='III', map_source=map_source_2), + ], + Reflector('丙丁甲乙', name='R', map_source=map_source_2), + rotate_up=True, rotate_after_type=True +) +assert e_custom_2.input('甲甲') == '乙丙' +e_custom_2.set_position() +assert e_custom_2.input('乙丙') == '甲甲' +``` + +### Plugboard + +Plugboard is also supported. + +```python +e.plugboard.plug('L', 'M') +e.plugboard.plug('O', 'P') +assert e.input('HELLOWORLDBYTHEAUTHOROFTHISPACKAGEDINGJUNYAO') == 'IMKPJAITPZIJNXSCSIJEOEDWZBRMCUOGQXGRJXQPFGHF' +e.set_position() +assert e.input('IMKPJAITPZIJNXSCSIJEOEDWZBRMCUOGQXGRJXQPFGHF') == 'HELLOWORLDBYTHEAUTHOROFTHISPACKAGEDINGJUNYAO' +e.unplug('L') +e.unplug('P') +``` + +## License + +MIT License. \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..f8d8975 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,6 @@ +[build-system] +requires = [ + "setuptools", + "wheel" +] +build-backend = "setuptools.build_meta" \ No newline at end of file diff --git a/setup.cfg b/setup.cfg new file mode 100644 index 0000000..b8b027c --- /dev/null +++ b/setup.cfg @@ -0,0 +1,28 @@ +[metadata] +name = enigma +version = 1.0.0 +author = DingJunyao +author_email = dingjunyao0703@163.com +description = Python-based Enigma. +long_description = file: README.md +long_description_content_type = text/markdown +url = https://github.com/DingJunyao/enigma +project_urls = + Source = https://github.com/DingJunyao/enigma + Bug Tracker = https://github.com/DingJunyao/enigma/issues +classifiers = + Programming Language :: Python :: 3 + License :: OSI Approved :: MIT License + Operating System :: OS Independent + Topic :: Security :: Cryptography + Topic :: Sociology :: History + Topic :: Education + Topic :: Utilities +[options] +package_dir = + = src +packages = find: +python_requires = >=3.6 + +[options.packages.find] +where = src \ No newline at end of file diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/enigma/__init__.py b/src/enigma/__init__.py new file mode 100644 index 0000000..0f20bc9 --- /dev/null +++ b/src/enigma/__init__.py @@ -0,0 +1 @@ +from .enigma import Enigma diff --git a/src/enigma/enigma.py b/src/enigma/enigma.py new file mode 100644 index 0000000..a79d2a2 --- /dev/null +++ b/src/enigma/enigma.py @@ -0,0 +1,100 @@ +from typing import Union + +from src.enigma.part import Plugboard +from src.enigma.part.plate import Rotor, EntryPlate, Reflector + + +class Enigma: + """An Enigma machine.""" + + def __init__( + self, rotor_list: Union[list[Rotor], tuple[Rotor]], reflector: Reflector, rotate_up: bool = False, + rotate_after_type: bool = False + ): + """Define the Enigma. + + :param rotor_list: The list of the rotors from right to left. Tuple with rotors is also supported. + :param reflector: The reflector of the Enigma. + :param rotate_up: Rotate direction of the rotors. + If it's True, it will rotate from small to large by index. + (If upper is A, lower is B, the plate will rotate from lower to upper when change from A to B.) + Or it's False (Default). + (If upper is B, lower is A, the plate will rotate from upper to lower when change from A to B.) + :param rotate_after_type: Set if the rotors rotate after type. Defaults to False. + :raise AttributeError: If the map_source of any rotor is not equal to that in reflector. + """ + self.rotors = list(rotor_list) + self.reflector = reflector + self.reflector.right_plate = self.rotors[-1] + self.entry_plate = EntryPlate(map_table=self.reflector.map_source, left_plate=self.rotors[0]) + for rotor_index, rotor in enumerate(self.rotors): + if rotor.map_source != self.reflector.map_source: + raise AttributeError(f'map_source of rotor[{rotor_index}] is not equal to that in reflector') + if rotor_index < len(self.rotors) - 1: + rotor.left_plate = self.rotors[rotor_index + 1] + else: + rotor.left_plate = self.reflector + if rotor_index > 0: + rotor.right_plate = self.rotors[rotor_index - 1] + else: + rotor.right_plate = self.entry_plate + for rotor in self.rotors: + rotor.rotate_up = rotate_up + self.rotate_after_type = rotate_after_type + self.plugboard = Plugboard(parent=self) + + def input(self, string: str) -> str: + """Input string to Enigma. + + :param string: String. + :return: String of encrypted result. + """ + new_string_list = [] + for letter_index, letter in enumerate(string): + if not self.rotate_after_type: + self.rotors[0].forward() + if letter in self.plugboard.map_dict: + letter = self.plugboard.map_dict[letter] + # print(f'Letter[{letter_index}]: {letter}') + for rotor_index, rotor in enumerate(self.rotors): + # print(rotor_index, 'i', letter) + # print(rotor_index, ' <', ''.join(rotor.current_state['right'])) + # print(rotor_index, '< ', ''.join(rotor.current_state['left'])) + letter = rotor.encrypt(letter, 'right') + # print(rotor_index, 'o', letter) + # print('R', 'i', letter) + # print('R', '< ', ''.join(self.reflector.current_state['left'])) + # print('R', '> ', ''.join(self.reflector.current_state['right'])) + letter = self.reflector.encrypt(letter, 'right') + # print('R', 'o', letter) + for rotor_index, rotor in enumerate(reversed(self.rotors)): + # print(rotor_index, 'i', letter) + # print(rotor_index, '> ', ''.join(rotor.current_state['left'])) + # print(rotor_index, ' >', ''.join(rotor.current_state['right'])) + letter = rotor.encrypt(letter, 'left') + # print(rotor_index, 'o', letter) + # print('I', 'i', letter) + # print('I', '< ', ''.join(self.entry_plate.current_state['left'])) + # print('I', ' >', ''.join(self.entry_plate.current_state['right'])) + letter = self.entry_plate.encrypt(letter, 'left') + # print('I', 'o', letter) + if letter in self.plugboard.map_dict: + letter = self.plugboard.map_dict[letter] + # print(f'Letter[{letter_index}]: {letter}') + new_string_list.append(letter) + if self.rotate_after_type: + self.rotors[0].forward() + return ''.join(new_string_list) + + def set_position(self, *position_list): + """Set positions of the rotors. + + :param position_list: the position of rotors from right to left. + :raise AttributeError: If position list length is not equal to rotors number. + """ + if not position_list: + position_list = (0, ) * len(self.rotors) + elif len(position_list) != len(self.rotors): + raise AttributeError('Position list length is not equal to rotors number') + for rotor, position in zip(self.rotors, position_list): + rotor.set_position(position) diff --git a/src/enigma/part/__init__.py b/src/enigma/part/__init__.py new file mode 100644 index 0000000..8f2bdcc --- /dev/null +++ b/src/enigma/part/__init__.py @@ -0,0 +1 @@ +from .plugboard import Plugboard diff --git a/src/enigma/part/plate/__init__.py b/src/enigma/part/plate/__init__.py new file mode 100644 index 0000000..8902fb8 --- /dev/null +++ b/src/enigma/part/plate/__init__.py @@ -0,0 +1,3 @@ +from .entry_plate import EntryPlate +from .rotor import Rotor +from .reflector import Reflector diff --git a/src/enigma/part/plate/entry_plate.py b/src/enigma/part/plate/entry_plate.py new file mode 100644 index 0000000..7d553ae --- /dev/null +++ b/src/enigma/part/plate/entry_plate.py @@ -0,0 +1,13 @@ +from typing import Union + +from .plate import Plate, _ALPHABET + + +class EntryPlate(Plate): + """Entry Plate of an Enigma.""" + + def __init__( + self, map_table: str = _ALPHABET, init_position: Union[str, int] = 0, left_plate: Plate = None, + name: str = 'Entry Plate' + ): + super().__init__(map_table, init_position, False, None, left_plate, name=name, map_source=map_table) diff --git a/src/enigma/part/plate/plate.py b/src/enigma/part/plate/plate.py new file mode 100644 index 0000000..fdf7555 --- /dev/null +++ b/src/enigma/part/plate/plate.py @@ -0,0 +1,138 @@ +from typing import Union, Optional + +_ALPHABET = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ' + + +class Plate: + """Plate of an Engima, including Entry Plate, Rotors (Walzen in German) and Reflector (Umkehrwalze in German).""" + + @staticmethod + def _map_table_check(map_table: str, map_source: str = _ALPHABET): + """Raise AttributeError if map table is invalid.""" + if len(set(map_source)) != len(map_source): + raise AttributeError(f'map_source "{map_source}" is invalid. It must be str with different letters.') + if len(set(map_table)) != len(map_source) or len(map_table) != len(map_source): + raise AttributeError( + f'map_table "{map_table}" is invalid. ' + f'It must be str with different letters, and every letter is in map_source "{map_source}".' + ) + for letter in map_source: + if letter not in map_table: + raise AttributeError(f'"{letter}" not in map table.') + + def __init_init_position_attr_exception(self, init_position: Union[int, str], map_source: str): + """Raise AttributeError if init_position is invalid.""" + raise AttributeError( + f'init_position "{init_position}" is invalid. ' + f'A letter string in map_source or an int between 0 and {len(map_source) - 1} is required.' + ) + + def __init__( + self, map_table: str = _ALPHABET, init_position: Union[str, int] = 0, auto_rotatable: bool = False, + right_plate=None, left_plate=None, rotate_up: bool = False, name: str = None, map_source: str = _ALPHABET + ): + """Defining a plate. + + :param map_table: String for map table. + :param init_position: Initial position. + :param auto_rotatable: If the plate is auto rotatable. + If it's True, the plate will be rotatable if other plate let it rotate. + Or it's False, but you can also let it rotate by setting its position. + :param right_plate: The plate on the right of this plate. + :param left_plate: The plate on the left of this plate. + :param rotate_up: Rotate direction of this plate. + If it's True, it will rotate from small to large by index. + (If upper is A, lower is B, the plate will rotate from lower to upper when change from A to B.) + Or it's False (Default). + (If upper is B, lower is A, the plate will rotate from upper to lower when change from A to B.) + :param name: The name of the plate. + :param map_source: String for map source. the circuits are: + item in map_source -> item in map_table + :raise AttributeError: If map_table or init_position is invalid. + """ + self._map_table_check(map_table, map_source) + dict_table = {} + for input_letter_index, input_letter in enumerate(map_source): + dict_table[input_letter] = map_table[input_letter_index] + self.dict_table = dict_table + if type(init_position) == str: + if len(init_position) != 1 or init_position not in map_source: + self.__init_init_position_attr_exception(init_position, map_source) + self.position = map_source.index(init_position) + elif type(init_position) == int: + if 0 <= init_position < len(map_table): + self.position = init_position + else: + self.__init_init_position_attr_exception(init_position, map_source) + else: + self.__init_init_position_attr_exception(init_position, map_source) + self.auto_rotatable = auto_rotatable + self.right_plate = right_plate + self.left_plate = left_plate + self.rotate_up = rotate_up + self.name = name + self.map_source = map_source + self.turnover = [] + + @property + def current_state(self) -> dict: + """Return a dict with 'left' and 'right' keys, which show the current state.""" + right = list(self.dict_table.keys()) + left = list(self.dict_table.values()) + transform_position = (-2 * self.rotate_up + 1) * self.position + current_right = right[transform_position:] + right[:transform_position] + current_left = left[transform_position:] + left[:transform_position] + return {'left': current_left, 'right': current_right} + + def transform_letter(self, letter: str, from_plate) -> str: + """Transform letter from previous process to the letter on the plate. + + :param letter: Letter from the previous process. + :param from_plate: 'left' or 'right' + :return: transformed letter. + """ + letter_index = from_plate.current_state['right'].index(letter) + return self.current_state['right'][letter_index] + + def encrypt(self, letter: str, letter_from: str = 'right') -> str: + """Return the output letter through the plate. + + :param letter: A letter. + :param letter_from: 'left' or 'right'. + :return: A letter. + :raise AttributeError: Input invaild + """ + if type(letter) != str or len(letter) != 1 or letter not in self.map_source: + raise AttributeError(f'Input "{letter}" is invalid. A letter string in map_source is required.') + if letter_from == 'left': + from_plate = self.left_plate + else: + from_plate = self.right_plate + if from_plate: + letter = self.transform_letter(letter, from_plate) + if letter_from == 'left': + key_index = list(self.dict_table.values()).index(letter) + left_letter = list(self.dict_table.keys())[key_index] + return left_letter + else: + return self.dict_table[letter] + + def set_position(self, position: Union[int, str]): + """Set position of the plate. + + :param position: Index or letter. + :raise AttributeError: If position is invalid. + """ + position_int: Optional[int] = None + if isinstance(position, str): + if position in self.map_source: + position_int = self.map_source.index(position) + elif isinstance(position, int): + if 0 <= position < len(self.map_source): + position_int = position + if position_int is None: + raise AttributeError( + f'position "{position}" is invaild. ' + f'A letter string in map_source or an int between 0 and {len(self.map_source) - 1} is required.' + ) + self.position = position_int diff --git a/src/enigma/part/plate/reflector.py b/src/enigma/part/plate/reflector.py new file mode 100644 index 0000000..33554b9 --- /dev/null +++ b/src/enigma/part/plate/reflector.py @@ -0,0 +1,23 @@ +from typing import Union + +from .plate import Plate, _ALPHABET + + +class Reflector(Plate): + """Reflector (Umkehrwalze in German) of an Enigma.""" + + def __init__( + self, map_table: str = _ALPHABET, init_position: Union[str, int] = 0, right_plate: Plate = None, + name: str = None, map_source: str = _ALPHABET + ): + """Defining a reflector. + + :param map_table: String for map table. + :param init_position: Initial position. + :param right_plate: The plate on the right of this plate. + :param name: The name of the plate. + :param map_source: String for map source. the circuits are: + item in map_source -> item in map_table + :raise AttributeError: If map_table or init_position is invalid. + """ + super().__init__(map_table, init_position, False, right_plate, None, name=name, map_source=map_source) diff --git a/src/enigma/part/plate/rotor.py b/src/enigma/part/plate/rotor.py new file mode 100644 index 0000000..0b331b4 --- /dev/null +++ b/src/enigma/part/plate/rotor.py @@ -0,0 +1,85 @@ +from typing import Union + +from .plate import Plate, _ALPHABET + + +class Rotor(Plate): + """Rotor (Walzen in German) of an Enigma.""" + + def __turnover_attr_exception( + self, turnover: Union[int, str, list[Union[int, str]], tuple[Union[int, str]]], map_source: str + ): + """Raise AttributeError if the turnover value is invalid.""" + raise AttributeError( + f'turnover "{turnover}" is invaild. ' + f'A letter string in map_source or an int between 0 and {len(map_source) - 1} is required.' + ) + + def turnover_single_check_and_standardize(self, turnover_single: Union[int, str], map_source: str) -> int: + """Check single turnover value and standardize it. + + :return: int: index of the letter. + :raise AttributeError: If the value is invalid. + """ + if isinstance(turnover_single, str): + if turnover_single in map_source: + return map_source.index(turnover_single) + elif isinstance(turnover_single, int): + if 0 <= turnover_single < len(map_source): + return turnover_single + self.__turnover_attr_exception(turnover_single, map_source) + + def __init__( + self, map_table: str = _ALPHABET, init_position: Union[str, int] = 0, auto_rotatable: bool = True, + right_plate: Plate = None, left_plate: Plate = None, rotate_up: bool = False, name: str = None, + map_source: str = _ALPHABET, turnover: Union[int, str, list[Union[int, str]], tuple[Union[int, str]]] = None + ): + """Defining a plate. + + :param map_table: String for map table. + :param init_position: Initial position. + :param auto_rotatable: If the plate is auto rotatable. + If it's True, the plate will be rotatable if other plate let it rotate. + Or it's False, but you can also let it rotate by setting its position. + :param right_plate: The plate on the right of this plate. + :param left_plate: The plate on the left of this plate. + :param rotate_up: Rotate direction of this plate. + If it's True, it will rotate from small to large by index. + (If upper is A, lower is B, the plate will rotate from lower to upper when change from A to B.) + Or it's False (Default). + (If upper is B, lower is A, the plate will rotate from upper to lower when change from A to B.) + :param name: The name of the plate. + :param map_source: String for map source. the circuits are: + item in map_source -> item in map_table + :param turnover: Define turnover position. + For example, if it's set to 'T', + it will let the left rotor rotate a step when this rotor rotates from T to U (in alphabet). + :raise AttributeError: If map_table or init_position is invalid. + """ + super().__init__(map_table, init_position, auto_rotatable, right_plate, left_plate, rotate_up, name, + map_source=map_source) + if turnover is None: + turnover = map_source[-1] + if isinstance(turnover, list) or isinstance(turnover, tuple): + turnover_list = [] + for turnover_item in turnover: + turnover_list.append(self.turnover_single_check_and_standardize(turnover_item, map_source)) + else: + turnover_list = [self.turnover_single_check_and_standardize(turnover, map_source)] + self.turnover = turnover_list + + def forward(self) -> int: + """Forward the plate. + + :return: Position of the plate. + """ + orignal_position = self.position + if orignal_position == len(self.map_source) - 1: + self.position = 0 + else: + self.position += 1 + if orignal_position in self.turnover: + if self.left_plate: + if self.left_plate.auto_rotatable: + self.left_plate.forward() + return self.position diff --git a/src/enigma/part/plugboard.py b/src/enigma/part/plugboard.py new file mode 100644 index 0000000..8254c55 --- /dev/null +++ b/src/enigma/part/plugboard.py @@ -0,0 +1,29 @@ +class Plugboard: + """Plugboard (Steckerbrett in German) of an Engima.""" + def __init__(self, parent=None): + self.map_dict = {} + self.parent = parent + + def check_parent(self): + if self.parent is None: + raise RuntimeError('parent of the Plugboard is None') + + def plug(self, letter_1, letter_2): + if letter_1 == letter_2: + raise AttributeError('Two letters can\'t be equal') + for letter in (letter_1, letter_2): + if letter not in self.parent.entry_plate.map_source: + raise AttributeError(f'"{letter}" is not in map_source') + if letter in self.map_dict and self.map_dict[letter]: + raise RuntimeError(f'"{letter}" have been set to be connected with another letter') + self.map_dict[letter_1] = letter_2 + self.map_dict[letter_2] = letter_1 + + def unplug(self, letter): + if letter not in self.parent.entry_plate.map_source: + raise AttributeError(f'"{letter}" is not in map_source') + if letter not in self.map_dict: + raise RuntimeError(f'"{letter}" have not been set') + another_letter = self.map_dict[letter] + del self.map_dict[letter] + del self.map_dict[another_letter] diff --git a/src/enigma/sample_plate/__init__.py b/src/enigma/sample_plate/__init__.py new file mode 100644 index 0000000..40617ef --- /dev/null +++ b/src/enigma/sample_plate/__init__.py @@ -0,0 +1,4 @@ +"""Source: +- http://users.telenet.be/d.rijmenants/Enigma%20Sim%20Manual.pdf +- http://www.ellsbury.com/ultraenigmawirings.htm +""" diff --git a/src/enigma/sample_plate/reflector.py b/src/enigma/sample_plate/reflector.py new file mode 100644 index 0000000..76b474c --- /dev/null +++ b/src/enigma/sample_plate/reflector.py @@ -0,0 +1,6 @@ +from src.enigma.part.plate import Reflector + +reflector_B = lambda: Reflector('YRUHQSLDPXNGOKMIEBFZCWVJAT', name='Reflector B') +reflector_C = lambda: Reflector('FVPJIAOYEDRZXWGCTKUQSBNMHL', name='Reflector C') +reflector_B_thin = lambda: Reflector('ENKQAUYWJICOPBLMDXZVFTHRGS', name='Reflector B Thin') +reflector_C_thin = lambda: Reflector('RDOBJNTKVEHMLFCWZAXGYIPSUQ', name='Reflector C Thin') diff --git a/src/enigma/sample_plate/rotor.py b/src/enigma/sample_plate/rotor.py new file mode 100644 index 0000000..d8889d6 --- /dev/null +++ b/src/enigma/sample_plate/rotor.py @@ -0,0 +1,12 @@ +from src.enigma.part.plate import Rotor + +rotor_I = lambda: Rotor('EKMFLGDQVZNTOWYHXUSPAIBRCJ', name='Rotor I', turnover='Q') +rotor_II = lambda: Rotor('AJDKSIRUXBLHWTMCQGZNPYFVOE', name='Rotor II', turnover='E') +rotor_III = lambda: Rotor('BDFHJLCPRTXVZNYEIWGAKMUSQO', name='Rotor III', turnover='V') +rotor_IV = lambda: Rotor('ESOVPZJAYQUIRHXLNFTGKDCMWB', name='Rotor IV', turnover='J') +rotor_V = lambda: Rotor('VZBRGITYUPSDNHLXAWMJQOFECK', name='Rotor V', turnover='Z') +rotor_VI = lambda: Rotor('JPGVOUMFYQBENHZRDKASXLICTW', name='Rotor VI', turnover=['Z', 'M']) +rotor_VII = lambda: Rotor('NZJHGRCXMYSWBOUFAIVLPEKQDT', name='Rotor VII', turnover=['Z', 'M']) +rotor_VIII = lambda: Rotor('FKQHTLXOCBJSPDZRAMEWNIUYGV', name='Rotor VIII', turnover=['Z', 'M']) +rotor_Beta = lambda: Rotor('LEYJVCNIXWPBQMDRTAKZGFUHOS', name='Rotor Beta') +rotor_Gamma = lambda: Rotor('FSOKANUERHMBTIYCWLQPZXVGJD', name='Rotor Gamma') diff --git a/tests/test_enigma.py b/tests/test_enigma.py new file mode 100644 index 0000000..5fe463b --- /dev/null +++ b/tests/test_enigma.py @@ -0,0 +1,70 @@ +from src.enigma import Enigma +from src.enigma.part.plate import Rotor, Reflector +from src.enigma.sample_plate.rotor import rotor_I, rotor_II, rotor_III +from src.enigma.sample_plate.reflector import reflector_B + + +def test_100_enigma(): + e = Enigma([rotor_III(), rotor_II(), rotor_I()], reflector_B()) + for input_letter, result in zip('ABCDEFGHIJKLMNOPQRSTUVWXYZ', 'BAQMFEXIHSWPDYTLCVJOZRKGNU'): + e.set_position() + assert e.input(input_letter) == result + e.set_position() + assert e.input('HELLOWORLDBYTHEAUTHOROFTHISPACKAGEDINGJUNYAO') == 'ILBDAAMTAZIJNXSCSIJJPDDWZBRCCUPGQXGRJXQOFGHL' + e.set_position() + assert e.input('ILBDAAMTAZIJNXSCSIJJPDDWZBRCCUPGQXGRJXQOFGHL') == 'HELLOWORLDBYTHEAUTHOROFTHISPACKAGEDINGJUNYAO' + e.set_position(0, 0, 0) + assert e.input('HELLOWORLDBYTHEAUTHOROFTHISPACKAGEDINGJUNYAO') == 'ILBDAAMTAZIJNXSCSIJJPDDWZBRCCUPGQXGRJXQOFGHL' + e.set_position(0, 0, 0) + assert e.input('ILBDAAMTAZIJNXSCSIJJPDDWZBRCCUPGQXGRJXQOFGHL') == 'HELLOWORLDBYTHEAUTHOROFTHISPACKAGEDINGJUNYAO' + e.set_position(25, 0, 0) + assert e.rotors[0].position == 25 + assert e.rotors[1].position == 0 + assert e.rotors[2].position == 0 + assert e.input('HELLOWORLDBYTHEAUTHOROFTHISPACKAGEDINGJUNYAO') == 'ZFEBMQKNGRKZEGBMJEQEVLCEVXFKHABTARMWPXSGALSX' + e.set_position(25, 0, 0) + assert e.input('ZFEBMQKNGRKZEGBMJEQEVLCEVXFKHABTARMWPXSGALSX') == 'HELLOWORLDBYTHEAUTHOROFTHISPACKAGEDINGJUNYAO' + e.set_position(3, 12, 21) + assert e.input('HELLOWORLDBYTHEAUTHOROFTHISPACKAGEDINGJUNYAO') == 'XTGHAGDIVUPGBZVQSFMBSGLKVQHQWESYRTSRMOOFGRLE' + e.set_position('D', 12, 21) + assert e.input('XTGHAGDIVUPGBZVQSFMBSGLKVQHQWESYRTSRMOOFGRLE') == 'HELLOWORLDBYTHEAUTHOROFTHISPACKAGEDINGJUNYAO' + + +def test_101_enigma_custom_source_1(): + map_source = 'ASDF' + e = Enigma( + [ + Rotor('AFSD', name='I', map_source=map_source), + Rotor('SDAF', name='II', map_source=map_source), + Rotor('DFAS', name='III', map_source=map_source), + ], + Reflector('DFAS', name='R', map_source=map_source), + rotate_up=True, rotate_after_type=True + ) + assert e.input('A') == 'S' + e.set_position() + assert e.input('S') == 'A' + e.set_position() + assert e.input('D') == 'F' + e.set_position() + assert e.input('F') == 'D' + e.set_position() + assert e.input('AA') == 'SD' + e.set_position() + assert e.input('SD') == 'AA' + + +def test_102_enigma_custom_source_2(): + map_source = '甲乙丙丁' + e = Enigma( + [ + Rotor('甲丁乙丙', name='I', map_source=map_source), + Rotor('乙丙甲丁', name='II', map_source=map_source), + Rotor('丙丁甲乙', name='III', map_source=map_source), + ], + Reflector('丙丁甲乙', name='R', map_source=map_source), + rotate_up=True, rotate_after_type=True + ) + assert e.input('甲甲') == '乙丙' + e.set_position() + assert e.input('乙丙') == '甲甲' diff --git a/tests/test_enigma_plug.py b/tests/test_enigma_plug.py new file mode 100644 index 0000000..576c466 --- /dev/null +++ b/tests/test_enigma_plug.py @@ -0,0 +1,50 @@ +from src.enigma import Enigma +from src.enigma.part.plate import Rotor, Reflector +from src.enigma.sample_plate.rotor import rotor_I, rotor_II, rotor_III +from src.enigma.sample_plate.reflector import reflector_B + + +def test_200_enigma_plug(): + e = Enigma([rotor_III(), rotor_II(), rotor_I()], reflector_B()) + e.plugboard.plug('L', 'M') + e.plugboard.plug('O', 'P') + for rotor in e.rotors: + assert rotor.position == 0 + for not_letter in 'ABCDEFGHIJKNQRSTUVWXYZ': + assert not_letter not in e.plugboard.map_dict + assert e.plugboard.map_dict['L'] == 'M' + assert e.plugboard.map_dict['M'] == 'L' + assert e.plugboard.map_dict['O'] == 'P' + assert e.plugboard.map_dict['P'] == 'O' + assert e.input('HELLOWORLDBYTHEAUTHOROFTHISPACKAGEDINGJUNYAO') == 'IMKPJAITPZIJNXSCSIJEOEDWZBRMCUOGQXGRJXQPFGHF' + e.set_position() + for rotor in e.rotors: + assert rotor.position == 0 + for not_letter in 'ABCDEFGHIJKNQRSTUVWXYZ': + assert not_letter not in e.plugboard.map_dict + assert e.plugboard.map_dict['L'] == 'M' + assert e.plugboard.map_dict['M'] == 'L' + assert e.plugboard.map_dict['O'] == 'P' + assert e.plugboard.map_dict['P'] == 'O' + assert e.input('IMKPJAITPZIJNXSCSIJEOEDWZBRMCUOGQXGRJXQPFGHF') == 'HELLOWORLDBYTHEAUTHOROFTHISPACKAGEDINGJUNYAO' + + +def test_201_enigma_plug_custom_source(): + map_source = 'ASDF' + e = Enigma( + [ + Rotor('AFSD', name='I', map_source=map_source), + Rotor('SDAF', name='II', map_source=map_source), + Rotor('DFAS', name='III', map_source=map_source), + ], + Reflector('DFAS', name='R', map_source=map_source), + rotate_up=True, rotate_after_type=True, + ) + e.plugboard.plug('A', 'D') + assert e.input('A') == 'F' + e.set_position() + assert e.input('S') == 'D' + e.set_position() + assert e.input('D') == 'S' + e.set_position() + assert e.input('F') == 'A' diff --git a/tests/test_plate.py b/tests/test_plate.py new file mode 100644 index 0000000..0942822 --- /dev/null +++ b/tests/test_plate.py @@ -0,0 +1,6 @@ +from src.enigma.part.plate.plate import Plate + + +def test_000_plate(): + r = Plate('EKMFLGDQVZNTOWYHXUSPAIBRCJ') + r.encrypt('A', 'right') \ No newline at end of file