enigma-cipher 0.1.0

//! An absurdly fast and highly flexible Enigma machine simulation, encryption, and decryption library.

mod alphabet;
mod reflector;
mod rotor;

use crate::alphabet::{Alphabet, AlphabetIndex, TryIntoAlphabetIndex, ALPHABET};
use crate::reflector::Reflector;
use crate::rotor::{Rotor, TryIntoRotors};

pub struct EnigmaMachine {
    rotors: (Rotor, Rotor, Rotor),
    ring_positions: (AlphabetIndex, AlphabetIndex, AlphabetIndex),
    ring_settings: (AlphabetIndex, AlphabetIndex, AlphabetIndex),
    reflector: Reflector,
    plugboard: std::collections::HashMap<char, char>,
}

impl EnigmaMachine {
    /// Creates a new Enigma machine with blank settings. The settings for the machine must be added using the methods
    /// of `EnigmaBuilder`; See the README for an example.
    ///
    /// The returned value from this will always be `Ok`, and will be an Enigma machine with rotors 1, 1, 1, ring positions
    /// 1, 1, and 1, ring settings 1, 1, and 1, reflector A, and an empty plugboard.
    #[allow(clippy::new_ret_no_self)]
    pub fn new() -> impl EnigmaBuilder {
        Ok(Self {
            rotors: (1, 1, 1).try_into_rotors().unwrap(),
            ring_positions: (1, 1, 1).try_into_alphabet_index().unwrap(),
            ring_settings: (1, 1, 1).try_into_alphabet_index().unwrap(),
            reflector: Reflector::A,
            plugboard: std::collections::HashMap::new(),
        })
    }

    /// Decodes the given text using this Enigma machine.
    ///
    /// This is exactly the same as calling `machine.encode(text)`, since the enigma cipher is
    /// symmetric; The only difference is semantic meaning and intent, i.e.,
    ///
    /// ```rust
    ///	assert_eq!(text, machine.decode(machine.decode(text)));
    ///	assert_eq!(text, machine.encode(machine.encode(text)));
    ///	assert_eq!(text, machine.decode(machine.encode(text)));
    ///	assert_eq!(text, machine.encode(machine.decode(text)));
    /// ```
    ///
    /// # Parameters
    /// - `text` - The text to decode.
    ///
    /// # Returns
    /// The decoded text.
    pub fn decode(&self, text: &str) -> String {
        let text = text.to_uppercase();
        let rotor_a = self.rotors.0.alphabet();
        let rotor_b = self.rotors.1.alphabet();
        let rotor_c = self.rotors.2.alphabet();

        let mut rotor_a_letter = self.ring_positions.0;
        let mut rotor_b_letter = self.ring_positions.1;
        let mut rotor_c_letter = self.ring_positions.2;

        let rotor_a_setting = self.ring_settings.0;
        let offset_a_setting = rotor_a_setting;
        let rotor_b_setting = self.ring_settings.1;
        let offset_b_setting = rotor_b_setting;
        let rotor_c_setting = self.ring_settings.2;
        let offset_c_setting = rotor_c_setting;

        let rotor_a = caeser_shift(&rotor_a.letters(), *offset_a_setting);
        let rotor_b = caeser_shift(&rotor_b.letters(), *offset_b_setting);
        let rotor_c = caeser_shift(&rotor_c.letters(), *offset_c_setting);

        let rotor_a_first_half = rotor_a.get((26 - *offset_a_setting as usize)..rotor_a.len()).unwrap().to_owned();
        let rotor_a_second_half = rotor_a.get(0..(26 - *offset_a_setting as usize)).unwrap().to_owned();
        let rotor_a = rotor_a_first_half + &rotor_a_second_half;
        let rotor_a = Alphabet::new(&rotor_a).unwrap();

        let rotor_b_first_half = rotor_b.get((26 - *offset_b_setting as usize)..rotor_b.len()).unwrap().to_owned();
        let rotor_b_second_half = rotor_b.get(0..(26 - *offset_b_setting as usize)).unwrap().to_owned();
        let rotor_b = rotor_b_first_half + &rotor_b_second_half;
        let rotor_b = Alphabet::new(&rotor_b).unwrap();

        let rotor_c_first_half = rotor_c.get((26 - *offset_c_setting as usize)..rotor_c.len()).unwrap().to_owned();
        let rotor_c_second_half = rotor_c.get(0..(26 - *offset_c_setting as usize)).unwrap().to_owned();
        let rotor_c = rotor_c_first_half + &rotor_c_second_half;
        let rotor_c = Alphabet::new(&rotor_c).unwrap();

        text.chars()
            .map(|mut letter| {
                // Non-alphabetic characters stay the same
                if !letter.is_alphabetic() {
                    return letter;
                }

                // Rotate rotor 3
                let mut rotor_trigger = self
                    .rotors
                    .2
                    .notches()
                    .iter()
                    .map(|notch| ALPHABET.index_of(*notch).unwrap())
                    .collect::<Vec<_>>()
                    .contains(&rotor_c_letter);
                rotor_c_letter += 1;

                // Rotate rotor 2
                if rotor_trigger {
                    rotor_trigger = self
                        .rotors
                        .1
                        .notches()
                        .iter()
                        .map(|notch| ALPHABET.index_of(*notch).unwrap())
                        .collect::<Vec<_>>()
                        .contains(&rotor_b_letter);
                    rotor_b_letter += 1;

                    // Rotate rotor 1
                    if rotor_trigger {
                        rotor_a_letter += 1;
                    }
                }
                // Double step sequence
                else if self
                    .rotors
                    .1
                    .notches()
                    .iter()
                    .map(|notch| ALPHABET.index_of(*notch).unwrap())
                    .collect::<Vec<_>>()
                    .contains(&rotor_b_letter)
                {
                    rotor_b_letter += 1;
                    rotor_a_letter += 1;
                }

                // Plugboard decryption
                if let Some(plugboarded_letter) = self.plugboard.get(&letter) {
                    letter = *plugboarded_letter;
                }

                let offset_a = rotor_a_letter;
                let offset_b = rotor_b_letter;
                let offset_c = rotor_c_letter;

                // Rotor 3 Encryption
                let pos = ALPHABET.index_of(letter).unwrap();
                let let_ = rotor_c.letter_at(pos + offset_c);
                let pos = ALPHABET.index_of(let_).unwrap();
                letter = ALPHABET.letter_at(pos - offset_c);

                // Rotor 2 Encryption
                let pos = ALPHABET.index_of(letter).unwrap();
                let let_ = rotor_b.letter_at(pos + offset_b);
                let pos = ALPHABET.index_of(let_).unwrap();
                letter = ALPHABET.letter_at(pos - offset_b);

                // Rotor 1 Encryption
                let pos = ALPHABET.index_of(letter).unwrap();
                let let_ = rotor_a.letter_at(pos + offset_a);
                let pos = ALPHABET.index_of(let_).unwrap();
                letter = ALPHABET.letter_at(pos - offset_a);

                // Reflector Encryption
                if let Some(reflected_letter) = self.reflector.alphabet().get(&letter) {
                    letter = *reflected_letter;
                }

                // Rotor 1 Encryption
                let pos = ALPHABET.index_of(letter).unwrap();
                let let_ = ALPHABET.letter_at(pos + offset_a);
                let pos = rotor_a.index_of(let_).unwrap();
                letter = ALPHABET.letter_at(pos - offset_a);

                // Rotor 2 Encryption
                let pos = ALPHABET.index_of(letter).unwrap();
                let let_ = ALPHABET.letter_at(pos + offset_b);
                let pos = rotor_b.index_of(let_).unwrap();
                letter = ALPHABET.letter_at(pos - offset_b);

                // Rotor 3 Encryption
                let pos = ALPHABET.index_of(letter).unwrap();
                let let_ = ALPHABET.letter_at(pos + offset_c);
                let pos = rotor_c.index_of(let_).unwrap();
                letter = ALPHABET.letter_at(pos - offset_c);

                // Plugboard Second Pass
                if let Some(plugboarded_letter) = self.plugboard.get(&letter) {
                    letter = *plugboarded_letter;
                }

                letter
            })
            .collect()
    }

    /// Encodes the given text using this Enigma machine.
    ///
    /// This is exactly the same as calling `machine.decode(text)`, since the enigma cipher is
    /// symmetric; The only difference is semantic meaning and intent, i.e.,
    ///
    /// ```rust
    ///	assert_eq!(text, machine.decode(machine.decode(text)));
    ///	assert_eq!(text, machine.encode(machine.encode(text)));
    ///	assert_eq!(text, machine.decode(machine.encode(text)));
    ///	assert_eq!(text, machine.encode(machine.decode(text)));
    /// ```
    ///
    /// # Parameters
    /// - `text` - The text to encode.
    ///
    /// # Returns
    /// The encoded text.
    pub fn encode(&self, text: &str) -> String {
        self.decode(text)
    }
}

/// A trait applied to `anyhow::Result<EnigmaMachine>` that allows building an enigma machine and passing along errors if they occur.
pub trait EnigmaBuilder {
    fn rotors(self, first: u8, second: u8, third: u8) -> anyhow::Result<EnigmaMachine>;

    /// Sets the plugboard for the machine. The given plugboard should be a space-separated string of letter pairs. This is automatically
    /// bidirectional, meaning the pair `AY` will map `A` to `Y` and also `Y` to `A`.
    ///
    /// # Parameters
    /// - `plugboard` - A space-separated string of letter pairs, i.e., `AY BF QR UX GZ`.
    ///
    /// # Returns
    /// The machine builder with the given plugboard applied.
    ///
    /// # Errors
    /// If the machine builder passed to this is already an error, an error is returned immediately.
    ///
    /// If the given plugboard contains duplicate letters, an error is returned.
    ///
    /// If the given plugboard is not formatted as a space-separated list of letter pairs, an error is returned.
    fn plugboard(self, plugboard: &str) -> anyhow::Result<EnigmaMachine>;

    fn reflector(self, reflector: &str) -> anyhow::Result<EnigmaMachine>;
    fn ring_settings(self, first: u8, second: u8, third: u8) -> anyhow::Result<EnigmaMachine>;

    /// Sets the "ring positions" or "rotor positions" of the machine.
    ///
    /// # Parameters
    /// - `first` - The offset of the first rotor, in `[1, 26]`.
    /// - `second` - The offset of the second rotor, in `[1, 26]`.
    /// - `third` - The offset of the third rotor, in `[1, 26]`.
    ///
    /// # Returns
    /// The machine builder with the given rotor positions applied.
    ///
    /// # Errors
    /// If the machine builder passed to this is already an error, an error is returned immediately.
    ///
    /// If the given numbers are not all in `[1, 26]`, an error is returned.
    fn ring_positions(self, first: u8, second: u8, third: u8) -> anyhow::Result<EnigmaMachine>;
}

impl EnigmaBuilder for anyhow::Result<EnigmaMachine> {
    fn ring_positions(self, first: u8, second: u8, third: u8) -> anyhow::Result<EnigmaMachine> {
        if let Ok(machine) = self {
            Ok(EnigmaMachine {
                ring_positions: (first, second, third).try_into_alphabet_index()?,
                ..machine
            })
        } else {
            self
        }
    }

    fn ring_settings(self, first: u8, second: u8, third: u8) -> anyhow::Result<EnigmaMachine> {
        if let Ok(machine) = self {
            Ok(EnigmaMachine {
                ring_settings: (first, second, third).try_into_alphabet_index()?,
                ..machine
            })
        } else {
            self
        }
    }

    fn reflector(self, reflector: &str) -> anyhow::Result<EnigmaMachine> {
        let reflector = Reflector::try_from(reflector)?;
        self.map(|machine| EnigmaMachine { reflector, ..machine })
    }

    fn rotors(self, first: u8, second: u8, third: u8) -> anyhow::Result<EnigmaMachine> {
        let rotors = (first, second, third).try_into_rotors()?;
        self.map(|machine| EnigmaMachine { rotors, ..machine })
    }

    fn plugboard(self, plugboard: &str) -> anyhow::Result<EnigmaMachine> {
        if let Ok(machine) = self {
            let mut chars = plugboard.chars().collect::<Vec<char>>();
            chars.dedup();
            if chars.len() != plugboard.len() {
                anyhow::bail!("Plugboard contains duplicate characters: {plugboard}");
            }

            let mappings = plugboard.split_whitespace();
            let mut plugboard = std::collections::HashMap::new();
            for pair in mappings {
                let mut chars = pair.chars();
                let first = chars.next().unwrap();
                let second = chars.next().unwrap();
                plugboard.insert(first, second);
                plugboard.insert(second, first);
            }

            Ok(EnigmaMachine { plugboard, ..machine })
        } else {
            self
        }
    }
}

fn caeser_shift(text: &str, amount: u8) -> String {
    text.chars()
        .map(|letter| {
            let code = letter as u8;
            if (65..=90).contains(&code) {
                (((code - 65 + amount) % 26) + 65) as char
            } else {
                letter
            }
        })
        .collect()
}