keyfork-mnemonic-generate/src/main.rs

@@ -3,13 +3,10 @@ use std::{
     error::Error,
     fs::{read_dir, read_to_string, File},
     io::Read,
-    fmt::Display,
 };
 
 use sha2::{Digest, Sha256};
 
-type Result<T, E = Box<dyn Error>> = std::result::Result<T, E>;
-
 /// Usage: keyfork-mnemonic-generate [bitsize: 128 or 256]
 /// CHECKS:
 /// * If the system is online
@@ -22,33 +19,6 @@ type Result<T, E = Box<dyn Error>> = std::result::Result<T, E>;
 static WARNING_LINKS: [&str; 1] =
     ["https://lore.kernel.org/lkml/20211223141113.1240679-2-Jason@zx2c4.com/"];
 
-#[derive(Debug)]
-enum MnemonicGenerationError {
-    InvalidByteCount(usize),
-    InvalidByteLength(usize),
-}
-
-impl MnemonicGenerationError {
-    fn boxed(self) -> Box<Self> {
-        Box::new(self)
-    }
-}
-
-impl Display for MnemonicGenerationError {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        match self {
-            MnemonicGenerationError::InvalidByteCount(count) => {
-                write!(f, "Invalid byte count: {count}, must be divisible by 8")
-            }
-            MnemonicGenerationError::InvalidByteLength(count) => {
-                write!(f, "Invalid byte length: {count}, must be 128 or 256")
-            }
-        }
-    }
-}
-
-impl Error for MnemonicGenerationError {}
-
 fn ensure_safe_kernel_version() {
     let kernel_version = read_to_string("/proc/version").expect("/proc/version");
     let v = kernel_version
@@ -91,64 +61,6 @@ fn ensure_offline() {
     }
 }
 
-// TODO: Can a Mnemonic be formatted using a wordlist, without allocating or without storing the
-// entire word list?
-struct Mnemonic {
-    words: Vec<usize>,
-    wordlist: Vec<String>,
-}
-
-impl Mnemonic {
-    fn from_entropy(bytes: &[u8]) -> Result<Mnemonic> {
-        let bit_count = bytes.len() * 8;
-        let hash = generate_slice_hash(bytes);
-
-        if bit_count % 32 != 0 {
-            return Err(MnemonicGenerationError::InvalidByteCount(bit_count).boxed());
-        }
-        // 192 supported for test suite
-        if ![128, 192, 256].contains(&bit_count) {
-            return Err(MnemonicGenerationError::InvalidByteLength(bit_count).boxed());
-        }
-        assert_eq!(bit_count % 32, 0, "bit count must be in 32 bit increments");
-
-        let mut bits = vec![false; bit_count + bit_count / 32];
-
-        for byte_index in 0..bit_count / 8 {
-            for bit_index in 0..8 {
-                bits[byte_index * 8 + bit_index] = (bytes[byte_index] & (1 << (7 - bit_index))) > 0;
-            }
-        }
-        for check_bit in 0..bit_count / 32 {
-            bits[bit_count + check_bit] = (hash[check_bit / 8] & (1 << (7 - (check_bit % 8)))) > 0;
-        }
-
-        assert_eq!(bits.len() % 11, 0, "unstable bit count");
-
-        let words = bits
-            .chunks_exact(11)
-            .map(|chunk| bitslice_to_usize(chunk.try_into().expect("11 bit chunks")))
-            .collect::<Vec<_>>();
-
-        let wordlist = build_wordlist();
-        Ok(Mnemonic { words, wordlist })
-    }
-}
-
-impl Display for Mnemonic {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        let mut iter = self.words.iter().peekable();
-        while let Some(word_index) = iter.next() {
-            let word = &self.wordlist[*word_index];
-            write!(f, "{word}")?;
-            if iter.peek().is_some() {
-                write!(f, " ")?;
-            }
-        }
-        Ok(())
-    }
-}
-
 fn generate_slice_hash(data: &[u8]) -> Vec<u8> {
     let mut hasher = Sha256::new();
     hasher.update(data);
@@ -172,7 +84,29 @@ fn bitslice_to_usize(bitslice: [bool; 11]) -> usize {
     index
 }
 
-fn main() -> Result<()> {
+fn entropy_to_bits(bytes: &[u8]) -> Vec<bool> {
+    let bit_count = bytes.len() * 8;
+    let hash = generate_slice_hash(bytes);
+
+    assert_eq!(bit_count % 32, 0, "bit count must be in 32 bit increments");
+
+    let mut bits = vec![false; bit_count + bit_count / 32];
+
+    for byte_index in 0..bit_count / 8 {
+        for bit_index in 0..8 {
+            bits[byte_index * 8 + bit_index] = (bytes[byte_index] & (1 << (7 - bit_index))) > 0;
+        }
+    }
+    for check_bit in 0..bit_count / 32 {
+        bits[bit_count + check_bit] = (hash[check_bit / 8] & (1 << (7 - (check_bit % 8)))) > 0;
+    }
+
+    assert_eq!(bits.len() % 11, 0, "unstable bit count");
+
+    bits
+}
+
+fn main() -> Result<(), Box<dyn Error>> {
     if !env::vars()
         .any(|(name, _)| name == "SHOOT_SELF_IN_FOOT" || name == "INSECURE_HARDWARE_ALLOWED")
     {
@@ -180,6 +114,13 @@ fn main() -> Result<()> {
         ensure_offline();
     }
 
+    let wordlist = build_wordlist();
+    assert_eq!(
+        wordlist.len(),
+        2usize.pow(11),
+        "Wordlist did not include correct word count"
+    );
+
     let bit_size: usize = env::args()
         .nth(1)
         .unwrap_or(String::from("256"))
@@ -194,9 +135,13 @@ fn main() -> Result<()> {
     let entropy = &mut [0u8; 256 / 8];
     random_handle.read_exact(&mut entropy[..])?;
 
-    let mnemonic = Mnemonic::from_entropy(&entropy[..bit_size / 8])?;
+    let seed_bits = entropy_to_bits(entropy);
 
-    println!("{mnemonic}");
+    let words = seed_bits
+        .chunks_exact(11)
+        .map(|chunk| wordlist[bitslice_to_usize(chunk.try_into().expect("11 bit chunks"))].clone())
+        .collect::<Vec<_>>();
+    println!("{}", words.join(" "));
 
     Ok(())
 }
@@ -216,9 +161,10 @@ mod tests {
     }
 
     #[test]
-    fn loads_mnemonics() -> Result<()> {
+    fn loads_mnemonics() -> Result<(), Box<dyn super::Error>> {
         let content = include_str!("test/vectors.json");
         let jsonobj: serde_json::Value = serde_json::from_str(content)?;
+        let wordlist = build_wordlist();
 
         for test in jsonobj["english"].as_array().unwrap() {
             let [ref hex_, ref seed, ..] = test.as_array().unwrap()[..] else {
@@ -226,9 +172,16 @@ mod tests {
             };
             let hex = hex::decode(hex_.as_str().unwrap()).unwrap();
 
-            let mnemonic = Mnemonic::from_entropy(&hex)?;
+            let seed_bits = entropy_to_bits(&hex);
 
-            assert_eq!(mnemonic.to_string(), seed.as_str().unwrap());
+            let words = seed_bits
+                .chunks_exact(11)
+                .map(|chunk| {
+                    wordlist[bitslice_to_usize(chunk.try_into().expect("11 bit chunks"))].clone()
+                })
+                .collect::<Vec<_>>();
+
+            assert_eq!(words.join(" "), seed.as_str().unwrap());
         }
         Ok(())
     }