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

use std::{collections::HashMap, str::FromStr, sync::Arc};

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 {
    /// Return an Arced version of the Wordlist
    #[allow(clippy::must_use_candidate)]
    pub fn arc(self) -> Arc<Self> {
        Arc::new(self)
    }

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

    fn inner(&self) -> &Vec<String> {
        &self.0
    }

    #[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 {
    words: Vec<usize>,
    wordlist: Arc<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.get_word(*word_index).expect("word");
            write!(f, "{word}")?;
            if iter.peek().is_some() {
                write!(f, " ")?;
            }
        }
        Ok(())
    }
}

/// The error type representing a failure to parse a [`Mnemonic`]. These errors only occur during
/// [`Mnemonic`] creation.
#[derive(Debug, Clone)]
pub enum MnemonicFromStrError {
    /// The amount of words used to parse a mnemonic was not correct.
    InvalidWordCount(usize),

    /// One of the words used to generate the mnemonic was not found in the default wordlist.
    InvalidWord(usize),
}

impl Display for MnemonicFromStrError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            MnemonicFromStrError::InvalidWordCount(count) => {
                write!(f, "Incorrect word count: {count}")
            }
            MnemonicFromStrError::InvalidWord(index) => {
                write!(f, "Unknown word at index: {index}")
            }
        }
    }
}

impl Error for MnemonicFromStrError {}

impl FromStr for Mnemonic {
    type Err = MnemonicFromStrError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let wordlist = Wordlist::default().arc();
        let hm: HashMap<&str, usize> = wordlist
            .inner()
            .iter()
            .enumerate()
            .map(|(a, b)| (b.as_str(), a))
            .collect();
        let mut words: Vec<usize> = Vec::with_capacity(24);
        for (index, word) in s.split_whitespace().enumerate() {
            match hm.get(&word) {
                Some(id) => words.push(*id),
                None => return Err(MnemonicFromStrError::InvalidWord(index)),
            }
        }
        if ![12, 24].contains(&words.len()) {
            return Err(MnemonicFromStrError::InvalidWordCount(words.len()));
        }
        Ok(Mnemonic { words, wordlist })
    }
}

fn generate_slice_hash(data: &[u8]) -> Vec<u8> {
    let mut hasher = Sha256::new();
    hasher.update(data);
    hasher.finalize().to_vec()
}

impl Mnemonic {
    /// 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: Arc<Wordlist>,
    ) -> Result<Mnemonic, 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>, Arc<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();
        let wordlist = Wordlist::default().arc();

        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 mnemonic = Mnemonic::from_entropy(&hex, wordlist.clone()).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().arc();
        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());
    }
}