rwh 0.11.3

Safe filetransfer through magic-wormhole
//! Wordlist generation and wormhole code utilities
use rand::seq::IndexedRandom;
use serde_json::{self, Value};
use std::fmt;

use super::Password;

/// Represents a list of words used to generate and complete wormhole codes.
/// A wormhole code is a sequence of words used for secure communication or identification.
#[derive(PartialEq)]
pub struct Wordlist {
	/// Number of words in a wormhole code
	num_words: usize,
	/// Odd and even wordlist
	words: Vec<Vec<String>>,
}

impl fmt::Debug for Wordlist {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		write!(f, "Wordlist ( {}, lots of words...)", self.num_words)
	}
}

impl Wordlist {
	#[cfg(test)]
	#[doc(hidden)]
	pub fn new(num_words: usize, words: Vec<Vec<String>>) -> Wordlist {
		Wordlist { num_words, words }
	}

	/// This function provides completion suggestions for a given `prefix` based on a word list.
	/// The completion method depends on the `fuzzy-complete` feature:
	/// - **With `fuzzy-complete` enabled**: Uses fuzzy search (approximate string matching) to find matches.
	/// - **Without `fuzzy-complete`**: Uses simple [`String::starts_with`] matching.
	pub fn get_completions(&self, prefix: &str) -> Vec<String> {
		let words = self.get_wordlist(prefix);

		let (prefix, partial) = match prefix.rsplit_once('-') {
			Some(t) => t,
			// Only channel number entered, don't return completions
			None => return Vec::new(),
		};

		#[cfg(feature = "fuzzy-complete")]
		let matches = self.fuzzy_complete(partial, &words);
		#[cfg(not(feature = "fuzzy-complete"))]
		let matches = self.normal_complete(partial, &words);

		matches.iter().map(|word| format!("{prefix}-{word}")).collect()
	}

	/// Counts `-` and returns correct wormhole code word list for completion
	fn get_wordlist<'a>(&'a self, prefix: &str) -> Vec<&'a str> {
		let count_dashes = prefix.matches('-').count();
		let index = 1 - (count_dashes % 2);
		self.words[index].iter().map(|w| w.as_str()).collect()
	}

	/// Fuzzy completes single wormhole code word
	#[cfg(feature = "fuzzy-complete")]
	fn fuzzy_complete<'a>(&self, partial: &str, words: &'a [&str]) -> Vec<&'a str> {
		// We use Jaro-Winkler algorithm because it emphasizes the beginning of a word
		use fuzzt::algorithms::JaroWinkler;

		fuzzt::get_top_n(partial, words, None, None, None, Some(&JaroWinkler))
	}

	/// Completes single wormhole code word using [`String::starts_with`]
	#[allow(dead_code)]
	fn normal_complete<'a>(&self, partial: &str, words: &'a [&str]) -> Vec<&'a str> {
		words.iter().filter(|word| !partial.is_empty() && word.starts_with(partial)).copied().collect()
	}

	/// Choose wormhole code word
	pub fn choose_words(&self) -> Password {
		let mut rng = rand::rng();

		let components = self
			.words
			.iter()
			.cycle()
			.take(self.num_words)
			.map(|words| words.choose(&mut rng).unwrap().as_str())
			.collect::<Vec<&str>>();

		#[expect(unsafe_code)]
		unsafe {
			Password::new_unchecked(components.join("-"))
		}
	}

	pub(crate) fn into_words(self) -> impl Iterator<Item = String> {
		self.words.into_iter().flatten()
	}

	/// Construct Wordlist struct with given number of words in a wormhole code
	pub fn default_wordlist(num_words: usize) -> Wordlist {
		Wordlist { num_words, words: load_pgpwords() }
	}
}

fn load_pgpwords() -> Vec<Vec<String>> {
	let raw_words_value: Value = serde_json::from_str(include_str!("pgpwords.json")).unwrap();
	let raw_words = raw_words_value.as_object().unwrap();
	let mut even_words: Vec<String> = Vec::with_capacity(256);
	even_words.resize(256, String::from(""));
	let mut odd_words: Vec<String> = Vec::with_capacity(256);
	odd_words.resize(256, String::from(""));
	for (index_str, values) in raw_words.iter() {
		let index = u8::from_str_radix(index_str, 16).unwrap() as usize;
		even_words[index] = values.get(1).unwrap().as_str().unwrap().to_lowercase().to_string();
		odd_words[index] = values.get(0).unwrap().as_str().unwrap().to_lowercase().to_string();
	}

	vec![even_words, odd_words]
}

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

	#[test]
	#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
	fn test_load_words() {
		let w = load_pgpwords();
		assert_eq!(w.len(), 2);
		assert_eq!(w[0][0], "adroitness");
		assert_eq!(w[1][0], "aardvark");
		assert_eq!(w[0][255], "yucatan");
		assert_eq!(w[1][255], "zulu");
	}

	#[test]
	#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
	fn test_default_wordlist() {
		let d = Wordlist::default_wordlist(2);
		assert_eq!(d.words.len(), 2);
		assert_eq!(d.words[0][0], "adroitness");
		assert_eq!(d.words[1][0], "aardvark");
		assert_eq!(d.words[0][255], "yucatan");
		assert_eq!(d.words[1][255], "zulu");
	}

	#[test]
	#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
	fn test_get_wordlist() {
		let list = Wordlist::default_wordlist(2);
		assert_eq!(list.words.len(), 2);
		assert_eq!(list.get_wordlist("22-"), &*list.words[0]);
		assert_eq!(list.get_wordlist("22-dictator-"), &*list.words[1]);
	}

	fn vec_strs(all: &str) -> Vec<&str> {
		all.split_whitespace().map(|s| if s == "." { "" } else { s }).collect()
	}

	fn vec_strings(all: &str) -> Vec<String> {
		vec_strs(all).iter().map(|s| (*s).to_owned()).collect()
	}

	#[test]
	#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
	fn test_completion() {
		let words: Vec<Vec<String>> = vec![vec_strings("purple green yellow"), vec_strings("sausages seltzer snobol")];

		let w = Wordlist::new(2, words);
		assert_eq!(w.get_completions(""), Vec::<String>::new());
		assert_eq!(w.get_completions("9"), Vec::<String>::new());
		assert_eq!(w.get_completions("seltz"), Vec::<String>::new());
		assert_eq!(w.get_completions("sausages-yello"), vec!["sausages-yellow".to_string()]);
	}

	#[test]
	#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
	fn test_choose_words() {
		let few_words: Vec<Vec<String>> = vec![vec_strings("purple"), vec_strings("sausages")];

		let w = Wordlist::new(2, few_words.clone());
		assert_eq!(w.choose_words().as_ref(), "purple-sausages");
		let w = Wordlist::new(3, few_words.clone());
		assert_eq!(w.choose_words().as_ref(), "purple-sausages-purple");
		let w = Wordlist::new(4, few_words);
		assert_eq!(w.choose_words().as_ref(), "purple-sausages-purple-sausages");
	}

	#[test]
	#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
	fn test_choose_words_matches_completion() {
		let few_words: Vec<Vec<String>> = vec![vec_strings("purple"), vec_strings("sausages")];

		let w = Wordlist::new(2, few_words.clone());
		assert_eq!(w.choose_words().as_ref(), "purple-sausages");

		// Check if odd and even wordlist are correctly selected
		assert_eq!(w.get_completions("1-purple-sausages").first().unwrap(), &format!("1-{}", w.choose_words().as_ref()));
	}

	#[test]
	#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
	fn test_choose_more_words() {
		let more_words = vec![vec_strings("purple yellow"), vec_strings("sausages")];

		let expected2 = vec_strs("purple-sausages yellow-sausages");
		let expected3 = ["purple-sausages-purple", "yellow-sausages-purple", "purple-sausages-yellow", "yellow-sausages-yellow"];

		let w = Wordlist::new(2, more_words.clone());
		for _ in 0..20 {
			assert!(expected2.contains(&w.choose_words().as_ref()));
		}

		let w = Wordlist::new(3, more_words);
		for _ in 0..20 {
			assert!(expected3.contains(&w.choose_words().as_ref()));
		}
	}

	#[test]
	#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
	#[cfg(feature = "fuzzy-complete")]
	fn test_completion_fuzzy() {
		let wl = Wordlist::default_wordlist(2);
		let list = wl.get_wordlist("22-");

		assert!(wl.fuzzy_complete("", &list).is_empty());

		assert_eq!(wl.fuzzy_complete("bzili", &list).first().unwrap(), &"brazilian");

		assert_eq!(wl.fuzzy_complete("carvan", &list).first().unwrap(), &"caravan");

		assert_ne!(wl.fuzzy_complete("choking", &list).first().unwrap(), &"choking")
	}

	#[test]
	#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
	fn test_completion_normal() {
		let wl = Wordlist::default_wordlist(2);
		let list = wl.get_wordlist("22-");

		assert!(wl.normal_complete("", &list).is_empty());

		assert_eq!(wl.normal_complete("braz", &list).first().unwrap(), &"brazilian");

		assert_eq!(wl.normal_complete("cara", &list).first().unwrap(), &"caravan");

		assert!(wl.normal_complete("cravan", &list).is_empty());
	}

	#[test]
	#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
	fn test_wormhole_code_normal_completions() {
		let list = Wordlist::default_wordlist(2);

		assert_eq!(list.get_completions("22"), Vec::<String>::new());
		assert_eq!(list.get_completions("22-"), Vec::<String>::new());
		#[cfg(not(feature = "fuzzy-complete"))]
		assert!(list.get_completions("22-troj").is_empty());

		assert_eq!(list.get_completions("22-compo").first().unwrap(), "22-component");

		assert_eq!(list.get_completions("22-component-check").first().unwrap(), "22-component-checkup");

		assert_ne!(list.get_completions("22-troj"), vec!["trojan".to_string()]);
	}

	#[test]
	#[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
	#[cfg(feature = "fuzzy-complete")]
	fn test_wormhole_code_fuzzy_completions() {
		let list = Wordlist::default_wordlist(2);

		assert_eq!(list.get_completions("22"), Vec::<String>::new());
		assert_eq!(list.get_completions("22-"), Vec::<String>::new());

		assert_eq!(list.get_completions("22-troj").first().unwrap(), "22-asteroid");

		assert_eq!(list.get_completions("22-decd").first().unwrap(), "22-decadence");

		assert_eq!(list.get_completions("22-udau").first().unwrap(), "22-undaunted");

		assert_eq!(list.get_completions("22-undua").first().unwrap(), "22-undaunted");

		assert_eq!(list.get_completions("22-undaunted-usht").first().unwrap(), "22-undaunted-upshot");
	}
}