1""" This is my second homework :)
 2     ,\
 3    # (_
 4      _[]##
 5  ###(())_
 6       ){}####
 7     _((
 8####/  )\
 9     ,;;`;,
10    (_______)
11      {===}
12      /===\
13     /=aat=\
14    The bonsai is PEP8 escape compatible
15    The bonsai is here because my ssd failed n times,
16        and the laptop is giving blue screen every 5 minutes(git files are also not physically saving,
17        just as the homework files in the git repo ...) :)
18        The homework function were having validation and exceptions, but yep in the past ....
19"""
20
21
22def plugboard(text, plugboard_map):
23    """
24    Plugboard for enigma machine
25
26    :param str text: The input text
27    :param list(two_element_sets()) plugboard_map: List of two element sets (Letter swap mapping)
28    :return: Encoded text
29    :rtype str
30    """
31    def plugboard_transform(letter):
32        for set_pair in plugboard_map:
33            if letter in set_pair:
34                return next(iter(set_pair.difference(set(letter))))
35        return letter
36
37    return ''.join([plugboard_transform(letter) for letter in text])
38
39
40def rotor(text, dictionary):
41    """
42    Rotor for enigma machine
43
44    :param str text: The input text
45    :param {one_letter_str,one_letter_str} dictionary: Rotor dictionary
46    :return: Encoded text
47    :rtype str
48    """
49    mapped_text = []
50    for letter in text:
51        if letter.isalpha():
52            mapped_text.append(dictionary.get(letter))
53        else:
54            mapped_text.append(letter)
55    return ''.join(mapped_text)
56
57
58def enigma_encrypt(plugboard_position, rotor_position):
59    """
60    Encryption function for enigma machine
61
62    :param list(two_element_set(str))) plugboard_position: Plugboard configuration
63    :param {one_letter_str,one_letter_str} rotor_position : Rotor dictionary configuration
64    :return: Decorator acting over function(str) giving the function enigma-encoded str
65    :rtype function-decorator
66    """
67    def decorator(func):
68        def encoder(text):
69            return func(rotor(plugboard(text, plugboard_position), rotor_position))
70        return encoder
71    return decorator
72
73
74def enigma_decrypt(plugboard_position, rotor_position):
75    """
76    Decryption function for enigma machine
77
78    :param list(two_element_set(str))) plugboard_position: Plugboard configuration
79    :param {one_letter_str,one_letter_str} rotor_position : Rotor dictionary configuration
80    :return: Decorator acting over function(str) giving the function enigma-encoded str
81    :rtype function-decorator
82    """
83    def decorator(func):
84        def decoder(text):
85            # Do you see the snake? :P
86            return func(plugboard(rotor(text, dict(zip(rotor_position.values(),
87                                                       rotor_position.keys()))), plugboard_position))
88        return decoder
89    return decorator