cfxkey/
brain_recover.rs

1// Copyright 2015-2019 Parity Technologies (UK) Ltd.
2// This file is part of Parity Ethereum.
3
4// Parity Ethereum is free software: you can redistribute it and/or modify
5// it under the terms of the GNU General Public License as published by
6// the Free Software Foundation, either version 3 of the License, or
7// (at your option) any later version.
8
9// Parity Ethereum is distributed in the hope that it will be useful,
10// but WITHOUT ANY WARRANTY; without even the implied warranty of
11// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12// GNU General Public License for more details.
13
14// You should have received a copy of the GNU General Public License
15// along with Parity Ethereum.  If not, see <http://www.gnu.org/licenses/>.
16
17use log::{info, trace};
18use std::collections::HashSet;
19
20use edit_distance::edit_distance;
21use parity_wordlist;
22
23use super::{Address, Brain, KeyPairGenerator};
24
25/// Tries to find a phrase for address, given the number
26/// of expected words and a partial phrase.
27///
28/// Returns `None` if phrase couldn't be found.
29pub fn brain_recover(
30    address: &Address, known_phrase: &str, expected_words: usize,
31) -> Option<String> {
32    let it = PhrasesIterator::from_known_phrase(known_phrase, expected_words);
33    for phrase in it {
34        let keypair = Brain::new(phrase.clone())
35            .generate()
36            .expect("Brain wallets are infallible; qed");
37        trace!("Testing: {}, got: {:?}", phrase, keypair.address());
38        if &keypair.address() == address {
39            return Some(phrase);
40        }
41    }
42
43    None
44}
45
46fn generate_substitutions(word: &str) -> Vec<&'static str> {
47    let mut words = parity_wordlist::WORDS
48        .iter()
49        .cloned()
50        .map(|w| (edit_distance(w, word), w))
51        .collect::<Vec<_>>();
52    words.sort_by(|a, b| a.0.cmp(&b.0));
53
54    words.into_iter().map(|pair| pair.1).collect()
55}
56
57/// Iterator over possible
58pub struct PhrasesIterator {
59    words: Vec<Vec<&'static str>>,
60    combinations: u64,
61    indexes: Vec<usize>,
62    has_next: bool,
63}
64
65impl PhrasesIterator {
66    pub fn from_known_phrase(
67        known_phrase: &str, expected_words: usize,
68    ) -> Self {
69        let known_words = parity_wordlist::WORDS
70            .iter()
71            .cloned()
72            .collect::<HashSet<_>>();
73        let mut words = known_phrase.split(' ')
74			.map(|word| match known_words.get(word) {
75				None => {
76					info!("Invalid word '{}', looking for potential substitutions.", word);
77					let substitutions = generate_substitutions(word);
78					info!("Closest words: {:?}", &substitutions[..10]);
79					substitutions
80				},
81				Some(word) => vec![*word],
82			})
83		.collect::<Vec<_>>();
84
85        // add missing words
86        if words.len() < expected_words {
87            let to_add = expected_words - words.len();
88            info!("Number of words is insuficcient adding {} more.", to_add);
89            for _ in 0..to_add {
90                words.push(parity_wordlist::WORDS.to_vec());
91            }
92        }
93
94        // start searching
95        PhrasesIterator::new(words)
96    }
97
98    pub fn new(words: Vec<Vec<&'static str>>) -> Self {
99        let combinations =
100            words.iter().fold(1u64, |acc, x| acc * x.len() as u64);
101        let indexes = words.iter().map(|_| 0).collect();
102        info!("Starting to test {} possible combinations.", combinations);
103
104        PhrasesIterator {
105            words,
106            combinations,
107            indexes,
108            has_next: combinations > 0,
109        }
110    }
111
112    pub fn combinations(&self) -> u64 { self.combinations }
113
114    fn current(&self) -> String {
115        let mut s = self.words[0][self.indexes[0]].to_owned();
116        for i in 1..self.indexes.len() {
117            s.push(' ');
118            s.push_str(self.words[i][self.indexes[i]]);
119        }
120        s
121    }
122
123    fn next_index(&mut self) -> bool {
124        let mut pos = self.indexes.len();
125        while pos > 0 {
126            pos -= 1;
127            self.indexes[pos] += 1;
128            if self.indexes[pos] >= self.words[pos].len() {
129                self.indexes[pos] = 0;
130            } else {
131                return true;
132            }
133        }
134
135        false
136    }
137}
138
139impl Iterator for PhrasesIterator {
140    type Item = String;
141
142    fn next(&mut self) -> Option<String> {
143        if !self.has_next {
144            return None;
145        }
146
147        let phrase = self.current();
148        self.has_next = self.next_index();
149        Some(phrase)
150    }
151}
152
153#[cfg(test)]
154mod tests {
155    use super::PhrasesIterator;
156
157    #[test]
158    fn should_generate_possible_combinations() {
159        let mut it = PhrasesIterator::new(vec![
160            vec!["1", "2", "3"],
161            vec!["test"],
162            vec!["a", "b", "c"],
163        ]);
164
165        assert_eq!(it.combinations(), 9);
166        assert_eq!(it.next(), Some("1 test a".to_owned()));
167        assert_eq!(it.next(), Some("1 test b".to_owned()));
168        assert_eq!(it.next(), Some("1 test c".to_owned()));
169        assert_eq!(it.next(), Some("2 test a".to_owned()));
170        assert_eq!(it.next(), Some("2 test b".to_owned()));
171        assert_eq!(it.next(), Some("2 test c".to_owned()));
172        assert_eq!(it.next(), Some("3 test a".to_owned()));
173        assert_eq!(it.next(), Some("3 test b".to_owned()));
174        assert_eq!(it.next(), Some("3 test c".to_owned()));
175        assert_eq!(it.next(), None);
176    }
177}
cfxkey/brain_recover.rs

cfxkey/
brain_recover.rs
