keyfork/keyfork-mnemonic-util/src/lib.rs

use sha2::{Digest, Sha256};
use std::{error::Error, fmt::Display};

/// The error type representing a failure to create a [`Mnemonic`]. These errors only occur during
/// [`Mnemonic`] creation.
#[derive(Debug, Clone)]
pub enum MnemonicGenerationError {
    /// The amount of bits passed to a mnemonic must be divisible by 32.
    InvalidByteCount(usize),

    /// The length of a mnemonic in bits must be within the BIP-0039 range, and supported by the
    /// library. Currently, only 128, 192 (for testing purposes), and 256 are supported.
    InvalidByteLength(usize),
}

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 {}

/// A BIP-0039 compatible list of words.
#[derive(Debug, Clone)]
pub struct Wordlist(Vec<String>);

impl Default for Wordlist {
    /// Returns the English wordlist in the Bitcoin BIP-0039 specification.
    fn default() -> Self {
        // TODO: English is the only supported language.
        let wordlist_file = include_str!("data/wordlist.txt");
        Wordlist(
            wordlist_file
                .lines()
                // skip 1: comment at top of file to point to BIP-0039 source.
                .skip(1)
                .map(|x| x.trim().to_string())
                .collect(),
        )
    }
}

impl Wordlist {
    /// Given an index, get a word from the wordlist.
    fn get_word(&self, word: usize) -> Option<&String> {
        self.0.get(word)
    }

    #[cfg(test)]
    fn into_inner(self) -> Vec<String> {
        self.0
    }
}

/// A BIP-0039 mnemonic with reference to a [`Wordlist`].
#[derive(Debug, Clone)]
pub struct Mnemonic<'a> {
    words: Vec<usize>,
    wordlist: &'a Wordlist,
}

impl<'a> Display for Mnemonic<'a> {
    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.get_word(*word_index).expect("word");
            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);
    hasher.finalize().to_vec()
}

impl<'a> Mnemonic<'a> {
    /// Generate a [`Mnemonic`] from the provided entropy and [`Wordlist`].
    ///
    /// # Errors
    /// An error may be returned if the entropy is not within the acceptable lengths.
    pub fn from_entropy(
        bytes: &[u8],
        wordlist: &'a Wordlist,
    ) -> Result<Mnemonic<'a>, MnemonicGenerationError> {
        let bit_count = bytes.len() * 8;
        let hash = generate_slice_hash(bytes);

        if bit_count % 32 != 0 {
            return Err(MnemonicGenerationError::InvalidByteCount(bit_count));
        }
        // 192 supported for test suite
        if ![128, 192, 256].contains(&bit_count) {
            return Err(MnemonicGenerationError::InvalidByteLength(bit_count));
        }

        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;
        }

        let words = bits
            // NOTE: Tested with all approved variants. Always divisible by 11.
            .chunks_exact(11)
            .map(|chunk| {
                // NOTE: usize to use for indexing wordlist later
                let mut num = 0usize;
                for i in 0..11 {
                    num += usize::from(chunk[10 - i]) << i;
                }
                num
            })
            .collect::<Vec<_>>();

        Ok(Mnemonic { words, wordlist })
    }

    #[must_use]
    pub fn into_inner(self) -> (Vec<usize>, &'a Wordlist) {
        (self.words, self.wordlist)
    }
}

#[cfg(test)]
mod tests {
    use std::{fs::File, io::Read};

    use super::*;

    #[test]
    fn wordlist_word_count_correct() {
        let wordlist = Wordlist::default().into_inner();
        assert_eq!(
            wordlist.len(),
            2usize.pow(11),
            "Wordlist did not include correct word count"
        );
    }

    #[test]
    fn conforms_to_trezor_tests() {
        let content = include_str!("data/vectors.json");
        let jsonobj: serde_json::Value = serde_json::from_str(content).unwrap();

        for test in jsonobj["english"].as_array().unwrap() {
            let [ref hex_, ref seed, ..] = test.as_array().unwrap()[..] else {
                panic!("bad test: {test}");
            };
            let hex = hex::decode(hex_.as_str().unwrap()).unwrap();

            let wordlist = Wordlist::default();
            let mnemonic = Mnemonic::from_entropy(&hex, &wordlist).unwrap();

            assert_eq!(mnemonic.to_string(), seed.as_str().unwrap());
        }
    }

    #[test]
    fn matches_bip39_crate() {
        let mut random_handle = File::open("/dev/random").unwrap();
        let entropy = &mut [0u8; 256 / 8];
        random_handle.read_exact(&mut entropy[..]).unwrap();
        let wordlist = Wordlist::default();
        let my_mnemonic = super::Mnemonic::from_entropy(&entropy[..256 / 8], &wordlist).unwrap();
        let their_mnemonic = bip39::Mnemonic::from_entropy(&entropy[..256 / 8]).unwrap();
        assert_eq!(my_mnemonic.to_string(), their_mnemonic.to_string());
    }
}