mini_enigma/
state_machine.rs

1//! Enigma Finite State Machine
2
3use super::utils::DEFAULT_ROTOR_NUM;
4
5use super::{
6    components::{Plugboard, Reflector, Rotor, Wheel},
7    utils::Letter,
8};
9
10/// ZST representing an unset, but required, setting in the enigma machine FSM
11pub struct Unset;
12
13#[derive(Debug)]
14/// Settings for the Enigma machine. Equivalent to a modern day IV (Initialisation Vector)
15pub struct EnigmaSettings<const ROTOR_NUM: usize = DEFAULT_ROTOR_NUM> {
16    /// Rotors and their ordering
17    pub rotors: [Rotor; ROTOR_NUM],
18    /// Plugboard cabling
19    pub plugboard: Plugboard,
20    /// Selected reflector
21    pub reflector: Reflector,
22    /// Ring settings for each rotor
23    pub ringstellung: [u8; ROTOR_NUM],
24    /// Starting position for each rotor
25    pub grundstellung: [Letter; ROTOR_NUM],
26}
27
28#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash, Debug)]
29/// Enigma Finite State Machine
30///
31/// An incomplete instance (e.g. from `Enigma::default`) can be initialised later
32/// semi-dynamically with the compiler checking correct initialisation before use.
33///
34/// # Examples
35///
36/// ```
37/// # use mini_enigma::state_machine::Enigma;
38/// # use mini_enigma::utils::Letter;
39/// # use mini_enigma::components::{Wheel, ROTOR_I, ROTOR_II, ROTOR_III, Plugboard, REFLECTOR_A};
40/// let enigma = Enigma::default()
41///     .set_rotors(Wheel::new([ROTOR_I, ROTOR_II, ROTOR_III]))
42///     .set_plugboard(Plugboard::new(&[]));
43///
44/// let mut enigma = enigma.set_reflector(REFLECTOR_A);
45///
46/// enigma.encrypt_letter(Letter::A);
47///
48/// ```
49///
50/// The compiler can detect if one of the required initialisation steps is missing:
51/// ```compile_fail,E0599
52/// # use mini_enigma::state_machine::Enigma;
53/// # use mini_enigma::utils::Letter;
54/// # use mini_enigma::components::{Wheel, ROTOR_I, ROTOR_II, ROTOR_III, Plugboard, REFLECTOR_A};
55/// let mut enigma = Enigma::default()
56///     .set_rotors(Wheel::new([ROTOR_I, ROTOR_II, ROTOR_III]))
57///     .set_plugboard(Plugboard::new(&[]));
58///
59/// enigma.encrypt_letter(Letter::A); // encrypt can't be called until a reflector is set
60///
61/// ```
62pub struct Enigma<R, P, O, const ROTOR_NUM: usize = DEFAULT_ROTOR_NUM> {
63    rotors: R,
64    plugboard: P,
65    reflector: O,
66}
67
68impl Default for Enigma<Unset, Unset, Unset, DEFAULT_ROTOR_NUM> {
69    /// Creates a default Enigma state machine where all parameters are unset.
70    ///
71    /// This function will always create an M3 variant Enigma.
72    ///
73    /// The configuration can be set later in execution using the various setter functions.
74    fn default() -> Self {
75        Enigma {
76            rotors: Unset,
77            plugboard: Unset,
78            reflector: Unset,
79        }
80    }
81}
82
83impl Enigma<Unset, Unset, Unset> {
84    /// Create a fully setup Enigma FSM from the provided settings
85    ///
86    /// # Examples
87    /// ```
88    /// # use mini_enigma::state_machine::{Enigma, EnigmaSettings};
89    /// # use mini_enigma::utils::Letter;
90    /// # use mini_enigma::components::{Wheel, ROTOR_II, ROTOR_IV, ROTOR_V, Plugboard, REFLECTOR_B};
91    /// # use core::str::FromStr;
92    /// let mut enigma = Enigma::new(&EnigmaSettings {
93    ///     rotors: [ROTOR_II, ROTOR_IV, ROTOR_V],
94    ///     plugboard: Plugboard::from_str(
95    ///         "AV,BS,CG,DL,FU,HZ,IN,KM,OW,RX",
96    ///     ).unwrap(),
97    ///     reflector: REFLECTOR_B,
98    ///     ringstellung: [2, 21, 12],
99    ///     grundstellung: [Letter::B, Letter::L, Letter::A],
100    /// });
101    /// enigma.encrypt_letter(Letter::A);
102    /// ```
103    ///
104    /// Rust will attempt to determine whether to create an M3 or M4 (or another variant) based on the size of the rotor array.
105    /// However, in some cases, it may be necessary to specify this explicitly:
106    /// ```
107    /// # use mini_enigma::state_machine::{Enigma, EnigmaSettings};
108    /// # use mini_enigma::utils::{Letter, M4};
109    /// # use mini_enigma::components::{Wheel, ROTOR_II, ROTOR_III, ROTOR_IV, ROTOR_V, Plugboard, REFLECTOR_B_THIN};
110    /// # use core::str::FromStr;
111    ///
112    /// let settings = EnigmaSettings {
113    ///     rotors: [ROTOR_II, ROTOR_IV, ROTOR_V, ROTOR_III],
114    ///     plugboard: Plugboard::from_str(
115    ///         "AV,BS,CG,DL,FU,HZ,IN,KM,OW,RX",
116    ///     ).unwrap(),
117    ///     reflector: REFLECTOR_B_THIN,
118    ///     ringstellung: [2, 21, 12, 3],
119    ///     grundstellung: [Letter::B, Letter::L, Letter::A, Letter::B],
120    /// };
121    ///
122    /// let mut enigma = Enigma::new::<M4>(&settings);
123    /// enigma.encrypt_letter(Letter::A);
124    /// ```
125    ///
126    /// As alluded to above this also lets you create instances of machines that never existed.
127    /// For example, a 10-rotor enigma machine can be created with `Enigma::new<10>(&settings)`
128    #[must_use]
129    pub fn new<const ROTOR_NUM: usize>(
130        config: &EnigmaSettings<ROTOR_NUM>,
131    ) -> Enigma<Wheel<ROTOR_NUM>, Plugboard, Reflector, ROTOR_NUM> {
132        let mut enigma = Enigma {
133            rotors: Wheel::new(config.rotors),
134            plugboard: config.plugboard,
135            reflector: config.reflector,
136        };
137        enigma.set_grundstellung(config.grundstellung);
138        enigma.set_ringstellung(config.ringstellung);
139        enigma
140    }
141}
142
143impl<AnyRotor, AnyPlugboard, AnyReflector, const ROTOR_NUM: usize>
144    Enigma<AnyRotor, AnyPlugboard, AnyReflector, ROTOR_NUM>
145{
146    /// Set the Enigma rotors
147    pub fn set_rotors(self, rotors: Wheel) -> Enigma<Wheel, AnyPlugboard, AnyReflector, ROTOR_NUM> {
148        Enigma {
149            rotors,
150            plugboard: self.plugboard,
151            reflector: self.reflector,
152        }
153    }
154
155    /// Set the Enigma plguboard
156    pub fn set_plugboard(
157        self,
158        plugboard: Plugboard,
159    ) -> Enigma<AnyRotor, Plugboard, AnyReflector, ROTOR_NUM> {
160        Enigma {
161            rotors: self.rotors,
162            plugboard,
163            reflector: self.reflector,
164        }
165    }
166
167    /// Set the Enigma reflector
168    pub fn set_reflector(
169        self,
170        reflector: Reflector,
171    ) -> Enigma<AnyRotor, AnyPlugboard, Reflector, ROTOR_NUM> {
172        Enigma {
173            rotors: self.rotors,
174            plugboard: self.plugboard,
175            reflector,
176        }
177    }
178}
179
180impl<AnyPlugboard, AnyReflector, const ROTOR_NUM: usize>
181    Enigma<Wheel<ROTOR_NUM>, AnyPlugboard, AnyReflector, ROTOR_NUM>
182{
183    /// Set the starting position for the rotors
184    pub fn set_grundstellung(
185        &mut self,
186        positions: [Letter; ROTOR_NUM],
187    ) -> &mut Enigma<Wheel<ROTOR_NUM>, AnyPlugboard, AnyReflector, ROTOR_NUM> {
188        self.rotors.set_positions(positions);
189        self
190    }
191
192    /// Set the ring settings for the rotors
193    pub fn set_ringstellung(
194        &mut self,
195        positions: [u8; ROTOR_NUM],
196    ) -> &mut Enigma<Wheel<ROTOR_NUM>, AnyPlugboard, AnyReflector, ROTOR_NUM> {
197        self.rotors.set_ring_wiring_offset(positions);
198        self
199    }
200
201    /// Get `Rotor` positions
202    pub fn get_grundstellung(&self) -> [Letter; ROTOR_NUM] {
203        self.rotors.get_positions()
204    }
205}
206
207impl<const ROTOR_NUM: usize> Enigma<Wheel<ROTOR_NUM>, Plugboard, Reflector, ROTOR_NUM> {
208    /// Encrypt/decrypt a `Letter` and advance the state
209    pub fn encrypt_letter(&mut self, letter: Letter) -> Letter {
210        self.rotors.step();
211        let left = self.rotors.forward(self.plugboard[letter]);
212        self.plugboard[self.rotors.backward(self.reflector[left])]
213    }
214
215    /// Encrypt/decrypt a `char` and advance the state
216    pub fn encrypt_char(&mut self, letter: char) -> char {
217        match Letter::try_from(letter) {
218            Ok(a) => self.encrypt_letter(a).into(),
219            Err(_) => letter,
220        }
221    }
222}
223
224#[cfg(test)]
225mod tests {
226    use super::super::components::{
227        REFLECTOR_A, REFLECTOR_B, ROTOR_I, ROTOR_II, ROTOR_III, ROTOR_IV, ROTOR_V,
228    };
229    use super::*;
230    use core::str::FromStr;
231
232    #[test]
233    fn test_machine() {
234        let enigma = Enigma::default();
235        let enigma = enigma.set_rotors(Wheel::new([ROTOR_I, ROTOR_II, ROTOR_III]));
236        let enigma = enigma.set_plugboard(Plugboard::new(&[]));
237        let enigma = enigma.set_reflector(REFLECTOR_A);
238        let mut enigma = enigma.set_rotors(Wheel::new([ROTOR_III, ROTOR_II, ROTOR_I]));
239        enigma.encrypt_letter(Letter::A);
240    }
241
242    #[test]
243    fn test_letter() {
244        let mut enigma = Enigma::default()
245            .set_rotors(Wheel::new([ROTOR_II, ROTOR_IV, ROTOR_V]))
246            .set_reflector(REFLECTOR_B)
247            .set_plugboard(Plugboard::from_str("AV,BS,CG,DL,FU,HZ,IN,KM,OW,RX").unwrap());
248        enigma
249            .set_grundstellung([Letter::B, Letter::L, Letter::A])
250            .set_ringstellung([2, 21, 12]);
251        assert_eq!(Letter::A, enigma.encrypt_letter(Letter::E));
252        assert_eq!(Letter::U, enigma.encrypt_letter(Letter::D));
253        assert_eq!(Letter::F, enigma.encrypt_letter(Letter::P));
254        assert_eq!(Letter::K, enigma.encrypt_letter(Letter::U));
255        assert_eq!(Letter::L, enigma.encrypt_letter(Letter::D));
256        assert_eq!(Letter::X, enigma.encrypt_letter(Letter::N));
257        assert_eq!(Letter::A, enigma.encrypt_letter(Letter::R));
258        assert_eq!(Letter::B, enigma.encrypt_letter(Letter::G));
259        assert_eq!(Letter::T, enigma.encrypt_letter(Letter::Y));
260        assert_eq!(Letter::E, enigma.encrypt_letter(Letter::S));
261    }
262
263    #[test]
264    fn test_string() {
265        let mut enigma = Enigma::new(&EnigmaSettings {
266            rotors: [ROTOR_II, ROTOR_IV, ROTOR_V],
267            plugboard: Plugboard::from_str("AV,BS,CG,DL,FU,HZ,IN,KM,OW,RX").unwrap(),
268            reflector: REFLECTOR_B,
269            ringstellung: [2, 21, 12],
270            grundstellung: [Letter::B, Letter::L, Letter::A],
271        });
272        let input = "AUFKLXABTE";
273        let expected = "EDPUDNRGYS".as_bytes();
274        for (i, ch) in input.chars().enumerate() {
275            assert_eq!(expected[i] as char, enigma.encrypt_char(ch));
276        }
277    }
278}
mini_enigma/state_machine.rs

mini_enigma/
state_machine.rs
