mini-enigma 0.3.0

//! Enigma Finite State Machine

use super::utils::DEFAULT_ROTOR_NUM;

use super::{
    components::{Plugboard, Reflector, Rotor, Wheel},
    utils::Letter,
};

/// ZST representing an unset, but required, setting in the enigma machine FSM
pub struct Unset;

#[derive(Debug)]
/// Settings for the Enigma machine. Equivalent to a modern day IV (Initialisation Vector)
pub struct EnigmaSettings<const ROTOR_NUM: usize = DEFAULT_ROTOR_NUM> {
    /// Rotors and their ordering
    pub rotors: [Rotor; ROTOR_NUM],
    /// Plugboard cabling
    pub plugboard: Plugboard,
    /// Selected reflector
    pub reflector: Reflector,
    /// Ring settings for each rotor
    pub ringstellung: [u8; ROTOR_NUM],
    /// Starting position for each rotor
    pub grundstellung: [Letter; ROTOR_NUM],
}

#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash, Debug)]
/// Enigma Finite State Machine
///
/// An incomplete instance (e.g. from `Enigma::default`) can be initialised later
/// semi-dynamically with the compiler checking correct initialisation before use.
///
/// # Examples
///
/// ```
/// # use mini_enigma::state_machine::Enigma;
/// # use mini_enigma::utils::Letter;
/// # use mini_enigma::components::{Wheel, ROTOR_I, ROTOR_II, ROTOR_III, Plugboard, REFLECTOR_A};
/// let enigma = Enigma::default()
///     .set_rotors(Wheel::new([ROTOR_I, ROTOR_II, ROTOR_III]))
///     .set_plugboard(Plugboard::new(&[]));
///
/// let mut enigma = enigma.set_reflector(REFLECTOR_A);
///
/// enigma.encrypt_letter(Letter::A);
///
/// ```
///
/// The compiler can detect if one of the required initialisation steps is missing:
/// ```compile_fail,E0599
/// # use mini_enigma::state_machine::Enigma;
/// # use mini_enigma::utils::Letter;
/// # use mini_enigma::components::{Wheel, ROTOR_I, ROTOR_II, ROTOR_III, Plugboard, REFLECTOR_A};
/// let mut enigma = Enigma::default()
///     .set_rotors(Wheel::new([ROTOR_I, ROTOR_II, ROTOR_III]))
///     .set_plugboard(Plugboard::new(&[]));
///
/// enigma.encrypt_letter(Letter::A); // encrypt can't be called until a reflector is set
///
/// ```
pub struct Enigma<R, P, O, const ROTOR_NUM: usize = DEFAULT_ROTOR_NUM> {
    rotors: R,
    plugboard: P,
    reflector: O,
}

impl Default for Enigma<Unset, Unset, Unset, DEFAULT_ROTOR_NUM> {
    /// Creates a default Enigma state machine where all parameters are unset.
    ///
    /// This function will always create an M3 variant Enigma.
    ///
    /// The configuration can be set later in execution using the various setter functions.
    fn default() -> Self {
        Enigma {
            rotors: Unset,
            plugboard: Unset,
            reflector: Unset,
        }
    }
}

impl Enigma<Unset, Unset, Unset> {
    /// Create a fully setup Enigma FSM from the provided settings
    ///
    /// # Examples
    /// ```
    /// # use mini_enigma::state_machine::{Enigma, EnigmaSettings};
    /// # use mini_enigma::utils::Letter;
    /// # use mini_enigma::components::{Wheel, ROTOR_II, ROTOR_IV, ROTOR_V, Plugboard, REFLECTOR_B};
    /// # use core::str::FromStr;
    /// let mut enigma = Enigma::new(&EnigmaSettings {
    ///     rotors: [ROTOR_II, ROTOR_IV, ROTOR_V],
    ///     plugboard: Plugboard::from_str(
    ///         "AV,BS,CG,DL,FU,HZ,IN,KM,OW,RX",
    ///     ).unwrap(),
    ///     reflector: REFLECTOR_B,
    ///     ringstellung: [2, 21, 12],
    ///     grundstellung: [Letter::B, Letter::L, Letter::A],
    /// });
    /// enigma.encrypt_letter(Letter::A);
    /// ```
    ///
    /// Rust will attempt to determine whether to create an M3 or M4 (or another variant) based on the size of the rotor array.
    /// However, in some cases, it may be necessary to specify this explicitly:
    /// ```
    /// # use mini_enigma::state_machine::{Enigma, EnigmaSettings};
    /// # use mini_enigma::utils::{Letter, M4};
    /// # use mini_enigma::components::{Wheel, ROTOR_II, ROTOR_III, ROTOR_IV, ROTOR_V, Plugboard, REFLECTOR_B_THIN};
    /// # use core::str::FromStr;
    ///
    /// let settings = EnigmaSettings {
    ///     rotors: [ROTOR_II, ROTOR_IV, ROTOR_V, ROTOR_III],
    ///     plugboard: Plugboard::from_str(
    ///         "AV,BS,CG,DL,FU,HZ,IN,KM,OW,RX",
    ///     ).unwrap(),
    ///     reflector: REFLECTOR_B_THIN,
    ///     ringstellung: [2, 21, 12, 3],
    ///     grundstellung: [Letter::B, Letter::L, Letter::A, Letter::B],
    /// };
    ///
    /// let mut enigma = Enigma::new::<M4>(&settings);
    /// enigma.encrypt_letter(Letter::A);
    /// ```
    ///
    /// As alluded to above this also lets you create instances of machines that never existed.
    /// For example, a 10-rotor enigma machine can be created with `Enigma::new<10>(&settings)`
    #[must_use]
    pub fn new<const ROTOR_NUM: usize>(
        config: &EnigmaSettings<ROTOR_NUM>,
    ) -> Enigma<Wheel<ROTOR_NUM>, Plugboard, Reflector, ROTOR_NUM> {
        let mut enigma = Enigma {
            rotors: Wheel::new(config.rotors),
            plugboard: config.plugboard,
            reflector: config.reflector,
        };
        enigma.set_grundstellung(config.grundstellung);
        enigma.set_ringstellung(config.ringstellung);
        enigma
    }
}

impl<AnyRotor, AnyPlugboard, AnyReflector, const ROTOR_NUM: usize>
    Enigma<AnyRotor, AnyPlugboard, AnyReflector, ROTOR_NUM>
{
    /// Set the Enigma rotors
    pub fn set_rotors(self, rotors: Wheel) -> Enigma<Wheel, AnyPlugboard, AnyReflector, ROTOR_NUM> {
        Enigma {
            rotors,
            plugboard: self.plugboard,
            reflector: self.reflector,
        }
    }

    /// Set the Enigma plguboard
    pub fn set_plugboard(
        self,
        plugboard: Plugboard,
    ) -> Enigma<AnyRotor, Plugboard, AnyReflector, ROTOR_NUM> {
        Enigma {
            rotors: self.rotors,
            plugboard,
            reflector: self.reflector,
        }
    }

    /// Set the Enigma reflector
    pub fn set_reflector(
        self,
        reflector: Reflector,
    ) -> Enigma<AnyRotor, AnyPlugboard, Reflector, ROTOR_NUM> {
        Enigma {
            rotors: self.rotors,
            plugboard: self.plugboard,
            reflector,
        }
    }
}

impl<AnyPlugboard, AnyReflector, const ROTOR_NUM: usize>
    Enigma<Wheel<ROTOR_NUM>, AnyPlugboard, AnyReflector, ROTOR_NUM>
{
    /// Set the starting position for the rotors
    pub fn set_grundstellung(
        &mut self,
        positions: [Letter; ROTOR_NUM],
    ) -> &mut Enigma<Wheel<ROTOR_NUM>, AnyPlugboard, AnyReflector, ROTOR_NUM> {
        self.rotors.set_positions(positions);
        self
    }

    /// Set the ring settings for the rotors
    pub fn set_ringstellung(
        &mut self,
        positions: [u8; ROTOR_NUM],
    ) -> &mut Enigma<Wheel<ROTOR_NUM>, AnyPlugboard, AnyReflector, ROTOR_NUM> {
        self.rotors.set_ring_wiring_offset(positions);
        self
    }

    /// Get `Rotor` positions
    pub fn get_grundstellung(&self) -> [Letter; ROTOR_NUM] {
        self.rotors.get_positions()
    }
}

impl<const ROTOR_NUM: usize> Enigma<Wheel<ROTOR_NUM>, Plugboard, Reflector, ROTOR_NUM> {
    /// Encrypt/decrypt a `Letter` and advance the state
    pub fn encrypt_letter(&mut self, letter: Letter) -> Letter {
        self.rotors.step();
        let left = self.rotors.forward(self.plugboard[letter]);
        self.plugboard[self.rotors.backward(self.reflector[left])]
    }

    /// Encrypt/decrypt a `char` and advance the state
    pub fn encrypt_char(&mut self, letter: char) -> char {
        match Letter::try_from(letter) {
            Ok(a) => self.encrypt_letter(a).into(),
            Err(_) => letter,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::super::components::{
        REFLECTOR_A, REFLECTOR_B, ROTOR_I, ROTOR_II, ROTOR_III, ROTOR_IV, ROTOR_V,
    };
    use super::*;
    use core::str::FromStr;

    #[test]
    fn test_machine() {
        let enigma = Enigma::default();
        let enigma = enigma.set_rotors(Wheel::new([ROTOR_I, ROTOR_II, ROTOR_III]));
        let enigma = enigma.set_plugboard(Plugboard::new(&[]));
        let enigma = enigma.set_reflector(REFLECTOR_A);
        let mut enigma = enigma.set_rotors(Wheel::new([ROTOR_III, ROTOR_II, ROTOR_I]));
        enigma.encrypt_letter(Letter::A);
    }

    #[test]
    fn test_letter() {
        let mut enigma = Enigma::default()
            .set_rotors(Wheel::new([ROTOR_II, ROTOR_IV, ROTOR_V]))
            .set_reflector(REFLECTOR_B)
            .set_plugboard(Plugboard::from_str("AV,BS,CG,DL,FU,HZ,IN,KM,OW,RX").unwrap());
        enigma
            .set_grundstellung([Letter::B, Letter::L, Letter::A])
            .set_ringstellung([2, 21, 12]);
        assert_eq!(Letter::A, enigma.encrypt_letter(Letter::E));
        assert_eq!(Letter::U, enigma.encrypt_letter(Letter::D));
        assert_eq!(Letter::F, enigma.encrypt_letter(Letter::P));
        assert_eq!(Letter::K, enigma.encrypt_letter(Letter::U));
        assert_eq!(Letter::L, enigma.encrypt_letter(Letter::D));
        assert_eq!(Letter::X, enigma.encrypt_letter(Letter::N));
        assert_eq!(Letter::A, enigma.encrypt_letter(Letter::R));
        assert_eq!(Letter::B, enigma.encrypt_letter(Letter::G));
        assert_eq!(Letter::T, enigma.encrypt_letter(Letter::Y));
        assert_eq!(Letter::E, enigma.encrypt_letter(Letter::S));
    }

    #[test]
    fn test_string() {
        let mut enigma = Enigma::new(&EnigmaSettings {
            rotors: [ROTOR_II, ROTOR_IV, ROTOR_V],
            plugboard: Plugboard::from_str("AV,BS,CG,DL,FU,HZ,IN,KM,OW,RX").unwrap(),
            reflector: REFLECTOR_B,
            ringstellung: [2, 21, 12],
            grundstellung: [Letter::B, Letter::L, Letter::A],
        });
        let input = "AUFKLXABTE";
        let expected = "EDPUDNRGYS".as_bytes();
        for (i, ch) in input.chars().enumerate() {
            assert_eq!(expected[i] as char, enigma.encrypt_char(ch));
        }
    }
}