<repo_name>pimvh/simpleenigma<file_sep>/parts/__init__.py
from .rotor import Rotor, StaticRotor
from .stekkerbrett import StekkerBrett<file_sep>/parts/stekkerbrett.py
from collections import OrderedDict
import itertools 
import random 
from .alphabetmapping import AlphabetMapping


class StekkerBrett(AlphabetMapping):

    def __init__(self):
        super().__init__()
        n_of_stekkers = random.randint(1, 9)

        self.mapping = OrderedDict(itertools.islice(self.mapping.items(), n_of_stekkers))
        self.reverse_mapping = OrderedDict(itertools.islice(self.reverse_mapping.items(), n_of_stekkers))
       
    def get_encoded_char(self, key) -> str:
        if not key in self.mapping:
            return key

        return self.mapping.get(key)
    

    def get_decoded_char(self, key) -> str:

        if not key in self.reverse_mapping:
            return key

        return self.reverse_mapping.get(key)
<file_sep>/parts/rotor.py
from collections import OrderedDict
from itertools import islice, cycle

from .alphabetmapping import AlphabetMapping


class Rotor(AlphabetMapping):
    """ defines a rotor of an enigma machine """
    def __init__(self) -> None:
        super().__init__()
        self.rotorsize = len(self.mapping)
        self.state = 0
        self.full_turns = 0
        self.shift = 1

        self.initial_state = next(iter(self.mapping))


    def get_encoded_char(self, key) -> str:
        return self.mapping.get(key)
    

    def get_decoded_char(self, key) -> str:
        return self.reverse_mapping.get(key)
    

    def get_initial_state(self) -> str:
        return self.initial_state


    def turn(self) -> bool:
        ''' turn the rotor 1 spot, return True when the rotor is fully turned '''
        self.state += self.shift

        self.mapping = OrderedDict(
            (k, v)
            for k, v in zip(self.mapping.keys(), islice(cycle(self.mapping.values()), 1, None))
        )

        self.reverse_mapping = OrderedDict((v, k) for k, v in self.mapping.items())

        encoded = self.mapping.get('a')
        # print(self.mapping.get('a'))
        decoded = self.reverse_mapping.get(encoded)
        # print(self.reverse_mapping.get('n'))
        assert decoded == 'a'

        if self.state > self.rotorsize:
            self.state = 0
            self.full_turns += 1
            return True
        
        return False
        

    def dump_settings(self):
        return self.mapping
    
    def set_settings(self, mapping):
        self.mapping = mapping


class StaticRotor(Rotor):
    """ defines a static rotor of an enigma machine """
    def __init__(self) -> None:
        super().__init__()
    
    def turn(self) -> bool:
        return False

    
<file_sep>/main.py
#!/usr/bin/env python3
from collections import OrderedDict
import copy
from typing import Text
import string
import time 

from parts import Rotor, StaticRotor, StekkerBrett

class Enigma():
    def __init__(self, rotors : list[Rotor], stekkerbret : StekkerBrett):        
        self.rotors = rotors
        self.stekkerbret = stekkerbret


    def get_state(self) -> list[OrderedDict]:
        return [rotor.dump_settings() for rotor in self.rotors]
    

    def set_state(self, settings) -> None:
        for rotor, setting in zip(self.rotors, settings):
            rotor.set_settings(setting)


    def encode(self, text_to_encode : Text) -> Text:
        print('encoding text...')

        result = ''        

        char_map = {' ' : 'x',
                    '?' : 'qq',
                    '.' : 'dot',
                    '/' : 'slash',
                    ',' : 'comma',}

        for c, char in enumerate(text_to_encode):
            
            if c % 4 == 0 and not c == 0:
                result += ' '

            if char in char_map:
                chars = char_map.get(char, '')
                
                for ch in chars:
                    result += self.encode_character(ch)
                continue

            char = char.lower()

            if not char in string.ascii_lowercase:
                continue
            
            result += self.encode_character(char)

        return result        
    
    def decode(self, text_to_decode : Text) -> Text:
        print('decoding text...')  

        result = ''

        for c, char in enumerate(text_to_decode.replace(' ', '')):
            result += self.decode_character(char)

        return result


    def encode_character(self, char : str):
        ''' parse a single character '''

        turn_next = True

        for i, current_rotor in enumerate(self.rotors):

            if turn_next:
                # print(f'rotor {i} turned')
                turn_next = current_rotor.turn()
                
            char = current_rotor.get_encoded_char(char)
        
        char = self.stekkerbret.get_encoded_char(char)

        return char
    
    def decode_character(self, char : str):
        ''' parse a single character '''

        char = self.stekkerbret.get_decoded_char(char)

        turn_next = True

        for i, current_rotor in enumerate(self.rotors):

            if turn_next:
                # print(f'rotor {i} turned')
                turn_next = current_rotor.turn()
                
        for i, current_rotor in enumerate(reversed(self.rotors)):
            char = current_rotor.get_decoded_char(char)
    
        return char

    def show_settings(self):
        pass

def main() -> None:

    rotors = []
    stekkerbrett = StekkerBrett()

    for i in range(10):
        rotors.append(Rotor())

    print(rotors[0])

    rotors.append(StaticRotor())

    decoder_rotors = [copy.deepcopy(x) for x in rotors]    
    decoder_stekkerbrett = copy.deepcopy(stekkerbrett)

    enigma = Enigma(rotors, stekkerbrett)
    enigma_decoder = Enigma(decoder_rotors, decoder_stekkerbrett)

    plaintext = '''Hoi dit is enigma'''
    print(f'{plaintext=}')

    encoded = enigma.encode(plaintext)
    print(f'{encoded=}')

    decoded = enigma_decoder.decode(encoded)

    print(f'{decoded=}')
    

if __name__ == "__main__":
    main()

<file_sep>/parts/alphabetmapping.py
from collections import OrderedDict
import random
import string

class AlphabetMapping():
    ''' implements a mapping of the alphabet for one letter to another '''

    def __init__(self) -> None:
        
        input_alpabet = [x for x in string.ascii_lowercase]
        random.shuffle(input_alpabet)
        output_alphabet = [x for x in string.ascii_lowercase]
        random.shuffle(output_alphabet)

        self.mapping : OrderedDict = OrderedDict()
        self.reverse_mapping : OrderedDict = OrderedDict()

        for key, value in zip(input_alpabet, output_alphabet):
            self.mapping.update({key : value})
            self.reverse_mapping.update({value : key})

        encoded = self.mapping.get('a')
        decoded = self.reverse_mapping.get(encoded)
        assert decoded == 'a'

    def show(self) -> None:
        ''' display the alphabet mapping of this wheel '''
        print(self)

    def __str__(self) -> str:
        out = ''
        for key, value in self.mapping.items():
            out += str(key) + ' --> ' + str(value) + '\n'
        
        return out