From 528dac825ae37caad6a1294e4672d5278ec5d64f Mon Sep 17 00:00:00 2001 From: a-moreira Date: Thu, 6 Jul 2023 13:39:59 -0300 Subject: [PATCH] re-use some code --- Cargo.lock | 1 + Cargo.toml | 1 + cli/src/main.rs | 47 +----------------------- src/core.rs | 2 +- src/core/wordlist.rs | 85 +++++++++++++++++++++++++++----------------- src/lib.rs | 4 ++- 6 files changed, 59 insertions(+), 81 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c05e8bb1..2a979b15 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1322,6 +1322,7 @@ dependencies = [ "base64 0.20.0", "bytecodec", "derive_more", + "dialoguer", "env_logger", "eyre", "futures", diff --git a/Cargo.toml b/Cargo.toml index 43869809..759b6043 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,6 +17,7 @@ edition = "2021" rust-version = "1.61" # MSRV (also change in CI) [dependencies] +dialoguer = { version = "0.10", features = ["completion"] } serde = { version = "1.0.120", features = ["rc"] } serde_json = "1.0.61" serde_derive = "1.0.120" diff --git a/cli/src/main.rs b/cli/src/main.rs index 773c9af9..612eb60b 100644 --- a/cli/src/main.rs +++ b/cli/src/main.rs @@ -15,6 +15,7 @@ use console::{style, Term}; use dialoguer::{Completion, Input}; use futures::{future::Either, Future, FutureExt}; use indicatif::{MultiProgress, ProgressBar}; +use magic_wormhole::PgpWordList; use magic_wormhole::{forwarding, transfer, transit, Wormhole}; use std::{ collections::HashMap, @@ -704,52 +705,6 @@ fn create_progress_bar(file_size: u64) -> ProgressBar { pb } -struct PgpWordList { - even_words: Vec, - odd_words: Vec, -} - -impl Default for PgpWordList { - fn default() -> Self { - let json = fs::read_to_string("./src/core/pgpwords.json").unwrap(); - let word_map: HashMap> = serde_json::from_str(&json).unwrap(); - let mut even_words: Vec = vec![]; - let mut odd_words: Vec = vec![]; - for (_idx, words) in word_map { - even_words.push(words[0].to_lowercase()); - odd_words.push(words[1].to_lowercase()); - } - - PgpWordList { - even_words, - odd_words, - } - } -} - -impl Completion for PgpWordList { - fn get(&self, input: &str) -> Option { - let count = input.matches("-").count(); - // we start with the odd words - let word_list: Vec = if count % 2 == 0 { - self.odd_words.clone() - } else { - self.even_words.clone() - }; - - let matches = word_list - .iter() - .filter(|word| word.starts_with(input)) - .collect::>(); - - if matches.len() == 1 { - Some(matches[0].to_string()) - } else { - None - } - } -} - fn enter_code() -> eyre::Result { let completion = PgpWordList::default(); let input = Input::new() diff --git a/src/core.rs b/src/core.rs index 27bf0d3c..dac9c9df 100644 --- a/src/core.rs +++ b/src/core.rs @@ -3,7 +3,7 @@ pub mod rendezvous; mod server_messages; #[cfg(test)] mod test; -mod wordlist; +pub mod wordlist; use serde_derive::{Deserialize, Serialize}; use std::borrow::Cow; diff --git a/src/core/wordlist.rs b/src/core/wordlist.rs index 01cc1f42..940d374a 100644 --- a/src/core/wordlist.rs +++ b/src/core/wordlist.rs @@ -2,31 +2,56 @@ use rand::{rngs::OsRng, seq::SliceRandom}; use serde_json::{self, Value}; use std::fmt; -#[derive(PartialEq)] -pub struct Wordlist { - pub num_words: usize, +use dialoguer::Completion; +use std::collections::HashMap; +use std::fs; + +pub struct PgpWordList { words: Vec>, + num_words: usize, } -impl fmt::Debug for Wordlist { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "Wordlist ( {}, lots of words...)", self.num_words) +impl PgpWordList { + pub fn choose_words(&self) -> String { + let mut rng = OsRng; + let components: Vec; + components = self + .words + .iter() + .cycle() + .take(self.num_words) + .map(|words| words.choose(&mut rng).unwrap().to_string()) + .collect(); + components.join("-") } } -impl Wordlist { - #[cfg(test)] - pub fn new(num_words: usize, words: Vec>) -> Wordlist { - Wordlist { num_words, words } +impl Default for PgpWordList { + fn default() -> Self { + let json = fs::read_to_string("./src/core/pgpwords.json").unwrap(); + let word_map: HashMap> = serde_json::from_str(&json).unwrap(); + let mut even_words: Vec = vec![]; + let mut odd_words: Vec = vec![]; + for (_idx, words) in word_map { + even_words.push(words[0].to_lowercase()); + odd_words.push(words[1].to_lowercase()); + } + let words = vec![even_words, odd_words]; + + PgpWordList { + words: words.clone(), + num_words: words.len(), + } } +} - #[allow(dead_code)] // TODO make this API public one day - pub fn get_completions(&self, prefix: &str) -> Vec { - let count_dashes = prefix.matches('-').count(); +impl Completion for PgpWordList { + fn get(&self, input: &str) -> Option { + let count_dashes = input.matches('-').count(); let mut completions = Vec::new(); let words = &self.words[count_dashes % self.words.len()]; - let last_partial_word = prefix.split('-').last(); + let last_partial_word = input.split('-').last(); let lp = if let Some(w) = last_partial_word { w.len() } else { @@ -34,12 +59,12 @@ impl Wordlist { }; for word in words { - let mut suffix: String = prefix.to_owned(); + let mut suffix: String = input.to_owned(); if word.starts_with(last_partial_word.unwrap()) { if lp == 0 { suffix.push_str(&word); } else { - let p = prefix.len() - lp; + let p = input.len() - lp; suffix.truncate(p as usize); suffix.push_str(&word); } @@ -51,21 +76,15 @@ impl Wordlist { completions.push(suffix); } } - completions.sort(); - completions - } - - pub fn choose_words(&self) -> String { - let mut rng = OsRng; - let components: Vec; - components = self - .words - .iter() - .cycle() - .take(self.num_words) - .map(|words| words.choose(&mut rng).unwrap().to_string()) - .collect(); - components.join("-") + if completions.len() == 1 { + Some(completions.first().unwrap().clone()) + } else if completions.len() == 0 { + None + } else { + // TODO: show vector of suggestions somehow + // println!("Suggestions: {:#?}", &completions); + None + } } } @@ -97,8 +116,8 @@ fn load_pgpwords() -> Vec> { vec![even_words, odd_words] } -pub fn default_wordlist(num_words: usize) -> Wordlist { - Wordlist { +pub fn default_wordlist(num_words: usize) -> PgpWordList { + PgpWordList { num_words, words: load_pgpwords(), } diff --git a/src/lib.rs b/src/lib.rs index 937b43c5..6b779fca 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -37,5 +37,7 @@ pub mod uri; pub use crate::core::{ key::{GenericKey, Key, KeyPurpose, WormholeKey}, - rendezvous, AppConfig, AppID, Code, Wormhole, WormholeError, WormholeWelcome, + rendezvous, + wordlist::PgpWordList, + AppConfig, AppID, Code, Wormhole, WormholeError, WormholeWelcome, };