enigma_cipher/
lib.rs

1//! An absurdly fast and highly flexible Enigma machine simulation, encryption, and decryption library.
2
3mod alphabet;
4mod reflector;
5mod rotor;
6
7use crate::alphabet::{Alphabet, AlphabetIndex, TryIntoAlphabetIndex, ALPHABET};
8use crate::reflector::Reflector;
9use crate::rotor::{Rotor, TryIntoRotors};
10
11pub type EnigmaResult<T> = anyhow::Result<T>;
12
13pub struct EnigmaMachine {
14    rotors: (Rotor, Rotor, Rotor),
15    ring_positions: (AlphabetIndex, AlphabetIndex, AlphabetIndex),
16    ring_settings: (AlphabetIndex, AlphabetIndex, AlphabetIndex),
17    reflector: Reflector,
18    plugboard: std::collections::HashMap<char, char>,
19}
20
21impl EnigmaMachine {
22    /// Creates a new Enigma machine with blank settings. The settings for the machine must be added using the methods
23    /// of `EnigmaBuilder`; See the README for an example.
24    ///
25    /// The returned value from this will always be `Ok`, and will be an Enigma machine with rotors 1, 1, 1, ring positions
26    /// 1, 1, and 1, ring settings 1, 1, and 1, reflector A, and an empty plugboard.
27    #[allow(clippy::new_ret_no_self)]
28    pub fn new() -> impl EnigmaBuilder {
29        Ok(Self {
30            rotors: (1, 1, 1).try_into_rotors().unwrap(),
31            ring_positions: (1, 1, 1).try_into_alphabet_index().unwrap(),
32            ring_settings: (1, 1, 1).try_into_alphabet_index().unwrap(),
33            reflector: Reflector::A,
34            plugboard: std::collections::HashMap::new(),
35        })
36    }
37
38    /// Decodes the given text using this Enigma machine.
39    ///
40    /// This is exactly the same as calling `machine.encode(text)`, since the enigma cipher is
41    /// symmetric; The only difference is semantic meaning and intent, i.e.,
42    ///
43    /// ```rust
44    ///	assert_eq!(text, machine.decode(machine.decode(text)));
45    ///	assert_eq!(text, machine.encode(machine.encode(text)));
46    ///	assert_eq!(text, machine.decode(machine.encode(text)));
47    ///	assert_eq!(text, machine.encode(machine.decode(text)));
48    /// ```
49    ///
50    /// # Parameters
51    /// - `text` - The text to decode.
52    ///
53    /// # Returns
54    /// The decoded text.
55    pub fn decode(&self, text: &str) -> String {
56        let text = text.to_uppercase();
57        let rotor_a = self.rotors.0.alphabet();
58        let rotor_b = self.rotors.1.alphabet();
59        let rotor_c = self.rotors.2.alphabet();
60
61        let mut rotor_a_letter = self.ring_positions.0;
62        let mut rotor_b_letter = self.ring_positions.1;
63        let mut rotor_c_letter = self.ring_positions.2;
64
65        let rotor_a_setting = self.ring_settings.0;
66        let offset_a_setting = rotor_a_setting;
67        let rotor_b_setting = self.ring_settings.1;
68        let offset_b_setting = rotor_b_setting;
69        let rotor_c_setting = self.ring_settings.2;
70        let offset_c_setting = rotor_c_setting;
71
72        let rotor_a = caeser_shift(&rotor_a.letters(), *offset_a_setting);
73        let rotor_b = caeser_shift(&rotor_b.letters(), *offset_b_setting);
74        let rotor_c = caeser_shift(&rotor_c.letters(), *offset_c_setting);
75
76        let rotor_a_first_half = rotor_a.get((26 - *offset_a_setting as usize)..rotor_a.len()).unwrap().to_owned();
77        let rotor_a_second_half = rotor_a.get(0..(26 - *offset_a_setting as usize)).unwrap().to_owned();
78        let rotor_a = rotor_a_first_half + &rotor_a_second_half;
79        let rotor_a = Alphabet::new(&rotor_a).unwrap();
80
81        let rotor_b_first_half = rotor_b.get((26 - *offset_b_setting as usize)..rotor_b.len()).unwrap().to_owned();
82        let rotor_b_second_half = rotor_b.get(0..(26 - *offset_b_setting as usize)).unwrap().to_owned();
83        let rotor_b = rotor_b_first_half + &rotor_b_second_half;
84        let rotor_b = Alphabet::new(&rotor_b).unwrap();
85
86        let rotor_c_first_half = rotor_c.get((26 - *offset_c_setting as usize)..rotor_c.len()).unwrap().to_owned();
87        let rotor_c_second_half = rotor_c.get(0..(26 - *offset_c_setting as usize)).unwrap().to_owned();
88        let rotor_c = rotor_c_first_half + &rotor_c_second_half;
89        let rotor_c = Alphabet::new(&rotor_c).unwrap();
90
91        text.chars()
92            .map(|mut letter| {
93                // Non-alphabetic characters stay the same
94                if !letter.is_alphabetic() {
95                    return letter;
96                }
97
98                // Rotate rotor 3
99                let mut rotor_trigger = self
100                    .rotors
101                    .2
102                    .notches()
103                    .iter()
104                    .map(|notch| ALPHABET.index_of(*notch).unwrap())
105                    .collect::<Vec<_>>()
106                    .contains(&rotor_c_letter);
107                rotor_c_letter += 1;
108
109                // Rotate rotor 2
110                if rotor_trigger {
111                    rotor_trigger = self
112                        .rotors
113                        .1
114                        .notches()
115                        .iter()
116                        .map(|notch| ALPHABET.index_of(*notch).unwrap())
117                        .collect::<Vec<_>>()
118                        .contains(&rotor_b_letter);
119                    rotor_b_letter += 1;
120
121                    // Rotate rotor 1
122                    if rotor_trigger {
123                        rotor_a_letter += 1;
124                    }
125                }
126                // Double step sequence
127                else if self
128                    .rotors
129                    .1
130                    .notches()
131                    .iter()
132                    .map(|notch| ALPHABET.index_of(*notch).unwrap())
133                    .collect::<Vec<_>>()
134                    .contains(&rotor_b_letter)
135                {
136                    rotor_b_letter += 1;
137                    rotor_a_letter += 1;
138                }
139
140                // Plugboard decryption
141                if let Some(plugboarded_letter) = self.plugboard.get(&letter) {
142                    letter = *plugboarded_letter;
143                }
144
145                let offset_a = rotor_a_letter;
146                let offset_b = rotor_b_letter;
147                let offset_c = rotor_c_letter;
148
149                // Rotor 3 Encryption
150                let pos = ALPHABET.index_of(letter).unwrap();
151                let let_ = rotor_c.letter_at(pos + offset_c);
152                let pos = ALPHABET.index_of(let_).unwrap();
153                letter = ALPHABET.letter_at(pos - offset_c);
154
155                // Rotor 2 Encryption
156                let pos = ALPHABET.index_of(letter).unwrap();
157                let let_ = rotor_b.letter_at(pos + offset_b);
158                let pos = ALPHABET.index_of(let_).unwrap();
159                letter = ALPHABET.letter_at(pos - offset_b);
160
161                // Rotor 1 Encryption
162                let pos = ALPHABET.index_of(letter).unwrap();
163                let let_ = rotor_a.letter_at(pos + offset_a);
164                let pos = ALPHABET.index_of(let_).unwrap();
165                letter = ALPHABET.letter_at(pos - offset_a);
166
167                // Reflector Encryption
168                if let Some(reflected_letter) = self.reflector.alphabet().get(&letter) {
169                    letter = *reflected_letter;
170                }
171
172                // Rotor 1 Encryption
173                let pos = ALPHABET.index_of(letter).unwrap();
174                let let_ = ALPHABET.letter_at(pos + offset_a);
175                let pos = rotor_a.index_of(let_).unwrap();
176                letter = ALPHABET.letter_at(pos - offset_a);
177
178                // Rotor 2 Encryption
179                let pos = ALPHABET.index_of(letter).unwrap();
180                let let_ = ALPHABET.letter_at(pos + offset_b);
181                let pos = rotor_b.index_of(let_).unwrap();
182                letter = ALPHABET.letter_at(pos - offset_b);
183
184                // Rotor 3 Encryption
185                let pos = ALPHABET.index_of(letter).unwrap();
186                let let_ = ALPHABET.letter_at(pos + offset_c);
187                let pos = rotor_c.index_of(let_).unwrap();
188                letter = ALPHABET.letter_at(pos - offset_c);
189
190                // Plugboard Second Pass
191                if let Some(plugboarded_letter) = self.plugboard.get(&letter) {
192                    letter = *plugboarded_letter;
193                }
194
195                letter
196            })
197            .collect()
198    }
199
200    /// Encodes the given text using this Enigma machine.
201    ///
202    /// This is exactly the same as calling `machine.decode(text)`, since the enigma cipher is
203    /// symmetric; The only difference is semantic meaning and intent, i.e.,
204    ///
205    /// ```rust
206    ///	assert_eq!(text, machine.decode(machine.decode(text)));
207    ///	assert_eq!(text, machine.encode(machine.encode(text)));
208    ///	assert_eq!(text, machine.decode(machine.encode(text)));
209    ///	assert_eq!(text, machine.encode(machine.decode(text)));
210    /// ```
211    ///
212    /// # Parameters
213    /// - `text` - The text to encode.
214    ///
215    /// # Returns
216    /// The encoded text.
217    pub fn encode(&self, text: &str) -> String {
218        self.decode(text)
219    }
220}
221
222/// A trait applied to `anyhow::Result<EnigmaMachine>` that allows building an enigma machine and passing along errors if they occur.
223pub trait EnigmaBuilder {
224    fn rotors(self, first: u8, second: u8, third: u8) -> anyhow::Result<EnigmaMachine>;
225
226    /// Sets the plugboard for the machine. The given plugboard should be a space-separated string of letter pairs. This is automatically
227    /// bidirectional, meaning the pair `AY` will map `A` to `Y` and also `Y` to `A`.
228    ///
229    /// # Parameters
230    /// - `plugboard` - A space-separated string of letter pairs, i.e., `AY BF QR UX GZ`.
231    ///
232    /// # Returns
233    /// The machine builder with the given plugboard applied.
234    ///
235    /// # Errors
236    /// If the machine builder passed to this is already an error, an error is returned immediately.
237    ///
238    /// If the given plugboard contains duplicate letters, an error is returned.
239    ///
240    /// If the given plugboard is not formatted as a space-separated list of letter pairs, an error is returned.
241    fn plugboard(self, plugboard: &str) -> anyhow::Result<EnigmaMachine>;
242
243    fn reflector(self, reflector: &str) -> anyhow::Result<EnigmaMachine>;
244    fn ring_settings(self, first: u8, second: u8, third: u8) -> anyhow::Result<EnigmaMachine>;
245
246    /// Sets the "ring positions" or "rotor positions" of the machine.
247    ///
248    /// # Parameters
249    /// - `first` - The offset of the first rotor, in `[1, 26]`.
250    /// - `second` - The offset of the second rotor, in `[1, 26]`.
251    /// - `third` - The offset of the third rotor, in `[1, 26]`.
252    ///
253    /// # Returns
254    /// The machine builder with the given rotor positions applied.
255    ///
256    /// # Errors
257    /// If the machine builder passed to this is already an error, an error is returned immediately.
258    ///
259    /// If the given numbers are not all in `[1, 26]`, an error is returned.
260    fn ring_positions(self, first: u8, second: u8, third: u8) -> anyhow::Result<EnigmaMachine>;
261}
262
263impl EnigmaBuilder for anyhow::Result<EnigmaMachine> {
264    fn ring_positions(self, first: u8, second: u8, third: u8) -> anyhow::Result<EnigmaMachine> {
265        if let Ok(machine) = self {
266            Ok(EnigmaMachine {
267                ring_positions: (first, second, third).try_into_alphabet_index()?,
268                ..machine
269            })
270        } else {
271            self
272        }
273    }
274
275    fn ring_settings(self, first: u8, second: u8, third: u8) -> anyhow::Result<EnigmaMachine> {
276        if let Ok(machine) = self {
277            Ok(EnigmaMachine {
278                ring_settings: (first, second, third).try_into_alphabet_index()?,
279                ..machine
280            })
281        } else {
282            self
283        }
284    }
285
286    fn reflector(self, reflector: &str) -> anyhow::Result<EnigmaMachine> {
287        let reflector = Reflector::try_from(reflector)?;
288        self.map(|machine| EnigmaMachine { reflector, ..machine })
289    }
290
291    fn rotors(self, first: u8, second: u8, third: u8) -> anyhow::Result<EnigmaMachine> {
292        let rotors = (first, second, third).try_into_rotors()?;
293        self.map(|machine| EnigmaMachine { rotors, ..machine })
294    }
295
296    fn plugboard(self, plugboard: &str) -> anyhow::Result<EnigmaMachine> {
297        if let Ok(machine) = self {
298            let mut chars = plugboard.chars().collect::<Vec<char>>();
299            chars.dedup();
300            if chars.len() != plugboard.len() {
301                anyhow::bail!("Plugboard contains duplicate characters: {plugboard}");
302            }
303
304            let mappings = plugboard.split_whitespace();
305            let mut plugboard = std::collections::HashMap::new();
306            for pair in mappings {
307                let mut chars = pair.chars();
308                let first = chars.next().unwrap();
309                let second = chars.next().unwrap();
310                plugboard.insert(first, second);
311                plugboard.insert(second, first);
312            }
313
314            Ok(EnigmaMachine { plugboard, ..machine })
315        } else {
316            self
317        }
318    }
319}
320
321fn caeser_shift(text: &str, amount: u8) -> String {
322    text.chars()
323        .map(|letter| {
324            let code = letter as u8;
325            if (65..=90).contains(&code) {
326                (((code - 65 + amount) % 26) + 65) as char
327            } else {
328                letter
329            }
330        })
331        .collect()
332}
enigma_cipher/lib.rs

enigma_cipher/
lib.rs
