enigma-machine/src/main.rs

mod enigma;
mod letter;
mod plugboard;
mod random;
mod reflector;
mod rotor;

pub use enigma::*;
pub use letter::*;
pub use plugboard::*;
pub use random::*;
pub use reflector::*;
pub use rotor::*;

use std::io;

fn main() {
    #[cfg(debug_assertions)]
    pretty_env_logger::init();

    let rotors = choose_rotors();
    let plugboard = choose_plugboard();
    let reflector = choose_reflector();
    println!("Plugboard configuration: {}", plugboard);
    let enigma_machine = Enigma::new(rotors, plugboard, reflector);

    loop {
        let mut enigma_machine = enigma_machine.clone();

        println!("Input some text to encrypt:");

        let input = clean_input(get_input_from_stdin());
        let output = enigma_machine.encode(&input);

        match output {
            Ok(output) => println!("Input: {}\nOutput: {}", input, output),
            Err(e) => println!("Wrong input: {}", e),
        }
    }
}

fn choose_rotors() -> Vec<Rotor> {
    let mut rotors = vec![];

    while rotors.is_empty() {
        rotors.clear();
        println!("Available rotors:");
        for rotor in Rotor::all() {
            println!("\t{}", rotor.variant);
        }
        println!("Choose rotors (comma-separated, right to left):");
        let input = clean_input(get_input_from_stdin());

        for choice in input.split(',') {
            let rotor = match choice {
                "I" => Rotor::variant(RotorVariant::I),
                "II" => Rotor::variant(RotorVariant::II),
                "III" => Rotor::variant(RotorVariant::III),
                "IV" => Rotor::variant(RotorVariant::IV),
                "V" => Rotor::variant(RotorVariant::V),
                "VI" => Rotor::variant(RotorVariant::VI),
                "VII" => Rotor::variant(RotorVariant::VII),
                "VIII" => Rotor::variant(RotorVariant::VIII),
                "IDENTITY" => Rotor::variant(RotorVariant::Identity),
                _ => {
                    println!("Invalid rotor choice: {}", choice);
                    break;
                }
            };

            rotors.push(rotor);
        }
    }

    rotors
}

fn choose_plugboard() -> Plugboard {
    let mut plugboard = None;

    while plugboard.is_none() {
        println!("Available plugboards:");
        println!("\tRandom");
        println!("\tIdentity");
        println!("Choose plugboard (only one):");
        let input = clean_input(get_input_from_stdin());

        plugboard = match &input[..] {
            "IDENTITY" => Some(Plugboard::identity()),
            "RANDOM" => Some(Plugboard::random()),
            _ => {
                println!("Invalid plugboard choice: {}", input);
                None
            }
        };
    }

    plugboard.unwrap()
}

fn choose_reflector() -> Reflector {
    let mut reflector = None;

    while reflector.is_none() {
        println!("Available reflectors:");
        for reflector in Reflector::all() {
            println!("\t{}", reflector.variant);
        }
        println!("Choose reflector (only one):");
        let input = clean_input(get_input_from_stdin());

        reflector = match &input[..] {
            "A" => Some(Reflector::variant(ReflectorVariant::A)),
            "B" => Some(Reflector::variant(ReflectorVariant::B)),
            "C" => Some(Reflector::variant(ReflectorVariant::C)),
            "BETA" => Some(Reflector::variant(ReflectorVariant::Beta)),
            "GAMMA" => Some(Reflector::variant(ReflectorVariant::Gamma)),
            "BTHIN" => Some(Reflector::variant(ReflectorVariant::BThin)),
            "CTHIN" => Some(Reflector::variant(ReflectorVariant::CThin)),
            "IDENTITY" => Some(Reflector::variant(ReflectorVariant::Identity)),
            _ => {
                println!("Invalid reflector choice: {}", input);
                None
            }
        };
    }

    reflector.unwrap()
}

fn get_input_from_stdin() -> String {
    let mut input = String::new();
    io::stdin().read_line(&mut input).unwrap();
    input
}

fn clean_input(input: String) -> String {
    input
        .trim()
        .chars()
        .filter(|c| c != &' ' && c != &'\r' && c != &'\n')
        .collect::<String>()
        .to_uppercase()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn enigma_machine_encodes_correctly() {
        let rotors = vec![
            Rotor::variant(RotorVariant::I),
            Rotor::variant(RotorVariant::II),
            Rotor::variant(RotorVariant::III),
        ];
        let plugboard = Plugboard::identity();
        let reflector = Reflector::variant(ReflectorVariant::B);
        let enigma_machine = Enigma::new(rotors, plugboard, reflector);

        let mut em1 = enigma_machine.clone();
        let mut em2 = enigma_machine;

        println!("Encoding string");
        let result = em1
            .encode("MAGNUSERPISSEDAARLIGTILLEAGUEOFLEGENDS")
            .unwrap();
        assert_eq!(result, "HTMPRVHFYZUHSANJWQCCAOHUTDRICUYUJJTHBE");

        println!("Decoding string");
        let result = em2
            .encode("HTMPRVHFYZUHSANJWQCCAOHUTDRICUYUJJTHBE")
            .unwrap();
        assert_eq!(result, "MAGNUSERPISSEDAARLIGTILLEAGUEOFLEGENDS");
    }

    #[test]
    fn random_plugboard_is_complete() {
        let plugboard = Plugboard::random();

        for letter in Letter::all() {
            println!("Testing if '{:?}' is mapped in plugboard", letter);
            plugboard.letter(letter);
        }
    }

    #[test]
    fn rotor_rotates_as_expected() {
        let mut rotor = Rotor::variant(RotorVariant::Identity);

        println!("{:?}", rotor.mappings);
        assert_eq!(rotor.forwards(Letter::A), Letter::A);
        rotor.rotate();
        assert_eq!(rotor.forwards(Letter::A), Letter::A);

        for _ in 0..26 {
            rotor.rotate();
        }

        println!("{:?}", rotor.mappings);
        assert_eq!(rotor.forwards(Letter::A), Letter::A);
    }
}