rusty-diceware

main.rs (11654B)

---

 /*
 rusty-diceware - a password / passphrasse generator using wordlists, with or without dice.
 Copyright (C) 2015-2022  Yuval Langer
 
 This program is free software: you can redistribute it and/or modify
 it under the terms of the GNU Affero General Public License as published by
 the Free Software Foundation, either version 3 of the License, or
 (at your option) any later version.
 
 This program is distributed in the hope that it will be useful,
 but WITHOUT ANY WARRANTY; without even the implied warranty of
 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 GNU Affero General Public License for more details.
 
 You should have received a copy of the GNU Affero General Public License
 along with this program.  If not, see <https://www.gnu.org/licenses/>.
  */
 
 extern crate getopts;
 extern crate rand;
 
 use std::env;
 use std::fs::File;
 use std::io::stdin;
 use std::io::Bytes;
 use std::io::Read;
 use std::process::exit;
 use std::str::FromStr;
 
 use crate::rand::Rng;
 use diceware_wordlists::Wordlist;
 use getopts::Options;
 use rand::rngs::ThreadRng;
 use rand::thread_rng;
 
 fn make_options() -> Options {
     let mut opts = Options::new();
     opts.optflag("h", "help", "This help message.");
     opts.optflag("e", "entropy", "Display number of entropy bits.");
     opts.optflag("r", "dicerolls", "Provide results of physical dice rolls. Word per line, same digit order as in the files, digits between and including 1 and 6.");
     opts.optopt("n", "nword", "Number of words in a passphrase.", "NWORD");
     opts.optopt(
         "d",
         "delimiter",
         "The delimiter character used to separate the words.",
         "DELIM",
     );
     opts.optopt("f", "wordlist-file", "Path to a wordlist file.", "FILE");
     opts.optopt(
         "l",
         "wordlist",
         "Wordlist to use. (efflong (default), effshort1, effshort2, minilock, reinhold, or beale)",
         "WORDLIST",
     );
     opts
 }
 
 fn print_usage(program: &str, opts: Options) {
     let brief = format!("Usage: {} [options]", program);
     print!("{}", opts.usage(&brief));
 }
 
 fn unknown_wordlist(wordlist_name: &str) -> ! {
     eprintln!(
         "Unknown wordlist: {}. Available wordlists: efflong (default), effshort1, effshort2, beale, reinhold, or minilock.",
         wordlist_name,
     );
     exit(-1)
 }
 
 struct ReadWordDieCasts<'a, T> {
     bytes_iter: &'a mut Bytes<T>,
     wordlist: &'a Vec<&'a str>,
 }
 
 impl<T> ReadWordDieCasts<'_, T> {
     fn new<'a>(
         bytes_iter: &'a mut Bytes<T>,
         wordlist: &'a Vec<&'a str>,
     ) -> ReadWordDieCasts<'a, T> {
         ReadWordDieCasts {
             bytes_iter,
             wordlist,
         }
     }
 }
 
 impl<T: Read> Iterator for ReadWordDieCasts<'_, T> {
     type Item = String;
 
     fn next(&mut self) -> Option<<ReadWordDieCasts<T> as Iterator>::Item> {
         read_single_word(self.bytes_iter, self.wordlist)
     }
 }
 
 fn read_single_word<T>(bytes_iter: &mut Bytes<T>, wordlist: &[&str]) -> Option<String>
 where
     T: Read,
 {
     let die_cast_per_word = die_cast_per_word(wordlist.len());
     let mut word_index = 0;
 
     for die_index in 0..=die_cast_per_word {
         // for die_cast_per_word == 5 we have:
         //
         // die_significance = die_cast_per_word - die_index
         // 5                = 5                 - 0 : just starting, may be an EOF, so next() == None and we then return None
         // 4                = 5                 - 1 .
         // 3                = 5                 - 2 .
         // 2                = 5                 - 3 .
         // 1                = 5                 - 4 .
         // 0                = 5                 - 5 : should be an end of line byte.
 
         if die_index == die_cast_per_word {
             let current_byte_result = match bytes_iter.next() {
                 Some(x) => x,
                 None => {
                     eprintln!("Each roll must end with a newline.");
                     exit(-1);
                 }
             };
 
             match current_byte_result {
                 Ok(b'\n') => {
                     return Some(wordlist[word_index].to_string());
                 }
                 Ok(n) => match n {
                     b'1' | b'2' | b'3' | b'4' | b'5' | b'6' => {
                         eprintln!(
                             "Too many rolls.  Each word must have exactly {} die rolls.",
                             die_cast_per_word
                         );
                         exit(-1);
                     }
                     err => {
                         eprintln!("Input must be either 1 to 6 inclusive digits or newline. (Got: ascii code {})", err);
                         exit(-1);
                     }
                 },
                 Err(err) => {
                     eprintln!("Unknown error: {}", err);
                     exit(-1);
                 }
             };
         } else {
             let current_byte_result = match bytes_iter.next() {
                 Some(x) => x,
                 None => {
                     if die_index == 0 {
                         // We are at a legal EOF state, right before a single word die casting sequence.
                         return None;
                     } else {
                         // We are in the middle of die casting and somehow stdin ended.
                         eprintln!(
                             "Bad number of rolls.  Each word must have exactly {} die rolls.",
                             die_cast_per_word
                         );
                         exit(-1);
                     }
                 }
             };
 
             let die_significance = die_cast_per_word - die_index - 1;
             word_index += match current_byte_result {
                 Ok(b'1') => 0,
                 Ok(b'2') => 6_usize.pow(die_significance),
                 Ok(b'3') => 2 * 6_usize.pow(die_significance),
                 Ok(b'4') => 3 * 6_usize.pow(die_significance),
                 Ok(b'5') => 4 * 6_usize.pow(die_significance),
                 Ok(b'6') => 5 * 6_usize.pow(die_significance),
                 Ok(b'\n') => {
                     eprintln!(
                         "Bad number of rolls.  Each word must have exactly {} die rolls.",
                         die_cast_per_word
                     );
                     exit(-1);
                 }
                 Ok(x) => {
                     eprintln!("Input must be either 1 to 6 inclusive digits or newline. (Got ascii code {})", x);
                     exit(-1);
                 }
                 Err(err) => {
                     eprintln!("Error while reading die casts: {}", err);
                     exit(-1);
                 }
             };
         };
     }
 
     Some(wordlist[word_index].to_string())
 }
 
 fn entropy(wordlist: &[&str]) -> f64 {
     (wordlist.len() as f64).log2()
 }
 
 fn entropyn(wordlist: Vec<&str>, n: u64) -> f64 {
     entropy(&wordlist) * (n as f64)
 }
 
 struct ReadRngWord<'a> {
     rng: &'a mut ThreadRng,
     wordlist: &'a Vec<&'a str>,
 }
 
 impl<'a> ReadRngWord<'a> {
     fn new(rng: &'a mut ThreadRng, wordlist: &'a Vec<&'a str>) -> ReadRngWord<'a> {
         ReadRngWord { rng, wordlist }
     }
 }
 
 impl<'a> Iterator for ReadRngWord<'a> {
     type Item = String;
 
     fn next(&mut self) -> Option<String> {
         Some(self.wordlist[self.rng.gen_range(0..self.wordlist.len())].to_string())
     }
 }
 
 fn load_wordlist_file(filepath: &str) -> String {
     let mut wordlist_file = match File::open(filepath) {
         Ok(ok) => ok,
         Err(err) => panic!("Unable to open file: {}; due to error: {}", filepath, err),
     };
     let mut wordlist_string = String::new();
     if let Err(err) = wordlist_file.read_to_string(&mut wordlist_string) {
         panic!("Unable to read file: {}; due to error: {}", filepath, err)
     }
     wordlist_string
 }
 
 fn die_cast_per_word(wordlist_length: usize) -> u32 {
     wordlist_length.ilog(6)
 }
 
 fn main() {
     let args: Vec<String> = env::args().collect();
     let program = &args[0];
 
     let opts = make_options();
 
     let matches = match opts.parse(&args[1..]) {
         Ok(m) => m,
         Err(f) => {
             println!("{}\n", f);
             print_usage(program, opts);
             exit(-1);
         }
     };
 
     if matches.opt_present("h") {
         print_usage(program, opts);
         return;
     };
 
     let is_physical_rolls = matches.opt_present("r");
 
     let word_num: u64 = matches
         .opt_str("n")
         .map_or(8, |n_str| n_str.parse::<u64>().ok().unwrap());
 
     let delimiter: char = matches
         .opt_str("d")
         .map_or(' ', |n_str| n_str.parse::<char>().ok().unwrap());
 
     let is_entropy_printed = matches.opt_present("entropy");
 
     let wordlist_name = if let Some(wordlist_option) = matches.opt_str("l") {
         match Wordlist::from_str(&wordlist_option.to_lowercase()) {
             Ok(list) => list,
             _ => unknown_wordlist(&wordlist_option),
         }
     } else {
         Wordlist::default()
     };
 
     let wordlist_string;
 
     let wordlist = if let Some(wordlist_filepath) = matches.opt_str("f") {
         wordlist_string = load_wordlist_file(&wordlist_filepath);
         wordlist_string
             .split('\n')
             .map(|x| x.trim())
             .filter(|x| x != &"")
             .collect::<Vec<&str>>()
     } else {
         wordlist_name.get_list().to_vec()
     };
     let mut words_count = 0;
 
     if is_physical_rolls {
         let stdin = stdin();
         let mut bytes_iter = stdin.bytes();
 
         let mut words = ReadWordDieCasts::new(&mut bytes_iter, &wordlist);
 
         if let Some(word) = words.next() {
             print!("{}", word);
 
             words_count = 1;
 
             for word in words {
                 print!("{}{}", delimiter, word);
                 words_count += 1;
             }
             println!();
         };
     } else if word_num != 0 {
         let mut rng = thread_rng();
 
         let mut words = ReadRngWord::new(&mut rng, &wordlist);
 
         print!("{}", words.next().unwrap());
 
         words_count = 1;
 
         for _ in 1..word_num {
             print!("{}{}", delimiter, words.next().unwrap());
             words_count += 1;
         }
 
         println!();
     } else {
         print_usage(program, opts);
         exit(-1);
     };
 
     if is_entropy_printed {
         println!("{}", entropyn(wordlist, words_count));
     }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
     use std::io::BufReader;
 
     #[test]
     fn test_read_single_word() {
         let wordlist = Wordlist::default().get_list().to_vec();
         let input = &b"11111\n"[..];
         let mut reader = BufReader::with_capacity(input.len(), input).bytes();
         assert_eq!(
             Some("abacus".to_string()),
             read_single_word(&mut reader, &wordlist)
         );
     }
 
     #[test]
     fn test_read_words_from_die_casting() {
         let wordlist = Wordlist::default().get_list().to_vec();
         let input = &b"11111\n66665\n"[..];
         let mut reader = BufReader::with_capacity(input.len(), input).bytes();
         let mut words_iter = ReadWordDieCasts::new(&mut reader, &wordlist);
         assert_eq!(Some("abacus".to_string()), words_iter.next());
         assert_eq!(Some("zoology".to_string()), words_iter.next());
         assert_eq!(None, words_iter.next());
     }
 
     // #[test]
     // fn test_read_words_from_rolls() {
     //     let wordlist = Wordlist::default().get_list().to_vec();
     //     let input = &b"11111\n11111\n"[..];
     //     let mut reader = std::io::BufReader::with_capacity(input.len(), input).bytes();
     //     assert_eq!("abacus", read_words_from_rolls::<std::io::Bytes<std::io::BufReader<&[u8]>>, u8>(reader, &wordlist).unwrap());
     // }
 }