1use std::str::FromStr;
2
3use bitcoin::bip32::DerivationPath;
4use getset::Getters;
5use regex::Regex;
6use serde::{Deserialize, Serialize};
7use tracing::{error, info};
8use zeroize::{Zeroize, ZeroizeOnDrop};
9
10use crate::error::RetrieverError;
11
12use super::exploration_step::{ExplorationStep, ExplorationStepHardness};
13
14#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Hash, Getters, Default)]
15#[get = "pub with_prefix"]
16pub struct ExplorationPath {
17 base_paths: Vec<DerivationPath>,
18 explore: Vec<ExplorationStep>,
19 depth: u32,
20 sweep: bool,
21}
22
23impl ExplorationPath {
24 pub fn new(
25 base_paths: Option<Vec<String>>,
26 explore_str: &str,
27 exploration_depth: u32,
28 sweep: bool,
29 ) -> Result<Self, RetrieverError> {
30 info!("Creation of exploration path started.");
31 let base_paths = match base_paths {
32 Some(base_paths) => base_paths
33 .iter()
34 .map(|base_path_string| {
35 DerivationPath::from_str(base_path_string)
36 .map_err(RetrieverError::from)
37 .unwrap()
38 })
39 .collect::<Vec<DerivationPath>>(),
40 None => vec![DerivationPath::from_str("m").unwrap()],
41 };
42 if !check_input_chars(explore_str) {
43 error!("Encountered invalid exploration path.");
44 return Err(RetrieverError::InvalidExplorationPath);
45 }
46 let explore_path_split = split_path_steps(explore_str);
47 if explore_path_split
48 .iter()
49 .any(|step| !check_step_sanity(step.clone()))
50 {
51 error!("Encountered invalid exploration path.");
52 return Err(RetrieverError::InvalidExplorationPath);
53 }
54
55 let mut explore = vec![];
56 for step in explore_path_split {
57 explore.push(translate_step_string_to_exploration_step(
58 step,
59 exploration_depth,
60 )?)
61 }
62 info!("Creation of exploration path finished successfully.");
63 Ok(ExplorationPath {
64 base_paths,
65 explore,
66 depth: exploration_depth,
67 sweep,
68 })
69 }
70
71 pub fn num_of_paths(&self) -> usize {
72 info!("Calculating the number of paths in exploration path.");
73 if self.explore.is_empty() {
74 0
75 } else {
76 self.base_paths.len()
77 * self
78 .explore
79 .iter()
80 .fold(1usize, |acc, step| acc * step.num_children() as usize)
81 }
82 }
83
84 pub fn num_of_paths_sweep(&self) -> usize {
85 info!("Calculating the number of sweep paths in exploration path.");
86 let mut num_paths = 1;
87 let sweep_exploration_paths = self.generate_sweep_exploration_paths();
88 for path in sweep_exploration_paths {
89 num_paths += path.num_of_paths()
90 }
91 num_paths
92 }
93
94 pub fn size(&self) -> usize {
95 if self.sweep {
96 self.num_of_paths_sweep()
97 } else {
98 self.num_of_paths()
99 }
100 }
101
102 pub fn generate_sweep_exploration_paths(&self) -> Vec<ExplorationPath> {
103 info!("Creating sweep exploration paths.");
104 let mut sweep_paths = vec![];
105 for i in 0..self.explore.len() + 1 {
106 sweep_paths.push(ExplorationPath {
107 explore: self.explore[..i].to_vec(),
108 depth: self.depth,
109 base_paths: self.base_paths.clone(),
110 sweep: self.sweep,
111 });
112 }
113 sweep_paths
114 }
115}
116
117impl Zeroize for ExplorationPath {
118 fn zeroize(&mut self) {
119 self.base_paths =
120 vec![DerivationPath::from_str("m/1024/1024/1204/1024").unwrap(); self.base_paths.len()];
121 self.explore.zeroize();
122 self.depth.zeroize();
123 self.sweep.zeroize();
124 }
125}
126
127impl ZeroizeOnDrop for ExplorationPath {}
128
129pub fn check_input_chars(input: &str) -> bool {
130 let regex = Regex::new(r"^[\d./'ha*]+$").unwrap();
131 regex.is_match(input)
132}
133
134pub fn split_path_steps(input: &str) -> Vec<String> {
135 let mut path_split = input.split('/').collect::<Vec<&str>>();
136 path_split.retain(|str| !str.is_empty());
137 path_split
138 .iter()
139 .map(|path_str| path_str.to_string())
140 .collect()
141}
142
143pub fn step_is_range(step: &str) -> bool {
144 let range_regex = Regex::new(r"^\d*(\.\.)?\d+[h'a]?$").unwrap();
145 range_regex.is_match(step)
146}
147
148pub fn step_is_wildcard(step: &str) -> bool {
149 let wildcard_regex = Regex::new(r"^\*[h'a]?$").unwrap();
150 wildcard_regex.is_match(step)
151}
152
153pub fn check_step_sanity(step: String) -> bool {
154 step_is_wildcard(&step) || step_is_range(&step)
155}
156
157pub fn extract_step_hardness(step: &str) -> ExplorationStepHardness {
158 match step.chars().last().unwrap() {
159 'h' | '\'' => ExplorationStepHardness::Hardened,
160 'a' => ExplorationStepHardness::HardenedAndNormal,
161 _ => ExplorationStepHardness::Normal,
162 }
163}
164
165pub fn translate_wildcard_step_string_to_exploration_step(
166 step_string: String,
167 exploration_depth: u32,
168) -> ExplorationStep {
169 let hardness = extract_step_hardness(&step_string);
170 let start_inclusive = 0;
171 let end_inclusive = exploration_depth;
172 ExplorationStep::new(start_inclusive, end_inclusive, hardness)
173}
174
175pub fn translate_range_step_string_to_exploration_step(
176 step_string: String,
177) -> Result<ExplorationStep, RetrieverError> {
178 let hardness = extract_step_hardness(&step_string);
179
180 let point_regex = Regex::new(r"^\d+[h'a]?$").unwrap();
181 let start_regex = Regex::new(r"^\d+\.\.").unwrap();
182 let end_regex = Regex::new(r"\.\.\d+").unwrap();
183
184 let start_inclusive = match point_regex.find(&step_string) {
185 Some(start) => start
186 .as_str()
187 .chars()
188 .filter(|char| char.is_ascii_digit())
189 .map(|char| char.to_string())
190 .collect::<Vec<String>>()
191 .join("")
192 .parse::<u32>()
193 .unwrap(),
194 None => match start_regex.find(&step_string) {
195 Some(start) => start
196 .as_str()
197 .chars()
198 .filter(|char| *char != '.')
199 .map(|char| char.to_string())
200 .collect::<Vec<String>>()
201 .join("")
202 .parse::<u32>()
203 .unwrap(),
204 None => 0u32,
205 },
206 };
207
208 let end_inclusive = match point_regex.find(&step_string) {
209 Some(end) => end
210 .as_str()
211 .chars()
212 .filter(|char| char.is_ascii_digit())
213 .map(|char| char.to_string())
214 .collect::<Vec<String>>()
215 .join("")
216 .parse::<u32>()
217 .unwrap(),
218 None => match end_regex.find(&step_string) {
219 Some(end) => end
220 .as_str()
221 .chars()
222 .filter(|char| *char != '.')
223 .map(|char| char.to_string())
224 .collect::<Vec<String>>()
225 .join("")
226 .parse::<u32>()
227 .unwrap(),
228 None => return Err(RetrieverError::InvalidStepRange),
229 },
230 };
231
232 if end_inclusive < start_inclusive {
233 return Err(RetrieverError::InvalidStepRange);
234 }
235
236 Ok(ExplorationStep::new(
237 start_inclusive,
238 end_inclusive,
239 hardness,
240 ))
241}
242
243pub fn translate_step_string_to_exploration_step(
244 step_string: String,
245 exploration_depth: u32,
246) -> Result<ExplorationStep, RetrieverError> {
247 if step_is_range(&step_string) {
248 Ok(translate_range_step_string_to_exploration_step(
249 step_string,
250 )?)
251 } else if step_is_wildcard(&step_string) {
252 Ok(translate_wildcard_step_string_to_exploration_step(
253 step_string,
254 exploration_depth,
255 ))
256 } else {
257 Err(RetrieverError::InvalidExplorationPath)
258 }
259}
260
261#[cfg(test)]
262mod tests {
263
264 use super::*;
265
266 #[test]
267 fn check_input_chars_works_01() {
268 assert!(check_input_chars("89/..90'/*"));
269 assert!(check_input_chars("*/*a/*'/*h"));
270 assert!(!check_input_chars("o"));
271 assert!(check_input_chars("*../90//09"));
272 assert!(!check_input_chars("+/*h/90/5"));
273 assert!(!check_input_chars("+/+"));
274 assert!(!check_input_chars("+/7"));
275 }
276
277 #[test]
278 fn split_path_steps_work_01() {
279 let str = "90/..9/*a";
280 let result = split_path_steps(str);
281 let expected = vec!["90".to_string(), "..9".to_string(), "*a".to_string()];
282 assert_eq!(result, expected);
283
284 let str = "90//..9/*a";
285 let result = split_path_steps(str);
286 let expected = vec!["90".to_string(), "..9".to_string(), "*a".to_string()];
287 assert_eq!(result, expected);
288
289 let str = "..89/4..9/*a/*'/*h/*/90..900h";
290 let result = split_path_steps(str);
291 let expected = vec![
292 "..89".to_string(),
293 "4..9".to_string(),
294 "*a".to_string(),
295 "*'".to_string(),
296 "*h".to_string(),
297 "*".to_string(),
298 "90..900h".to_string(),
299 ];
300 assert_eq!(result, expected);
301 }
302
303 #[test]
304 fn step_is_range_works_01() {
305 let range_steps = vec!["..90", "8..78", "..4h", "8..9'", "9..9a"];
306 range_steps
307 .iter()
308 .for_each(|step| assert!(step_is_range(step)));
309
310 let not_range_steps = vec![
311 "*", "*'", "*h", "*a", "p", "..*", "h..*", "*'ha", "..*h", "*ha", "89'h",
312 ];
313 not_range_steps
314 .iter()
315 .for_each(|step| assert!(!step_is_range(step)));
316 }
317
318 #[test]
319 fn step_is_wildcard_works_01() {
320 let not_wildcard_steps = vec![
321 "..90", "8..78", "..4h", "8..9'", "9..9a", "**", "..*h", "*ha", "89'h",
322 ];
323 not_wildcard_steps
324 .iter()
325 .for_each(|step| assert!(!step_is_wildcard(step)));
326
327 let wildcard_steps = vec!["*", "*'", "*h", "*a"];
328 wildcard_steps
329 .iter()
330 .for_each(|step| assert!(step_is_wildcard(step)));
331 }
332
333 #[test]
334 fn extract_step_hardness_works_01() {
335 assert_eq!(
336 extract_step_hardness("9'"),
337 ExplorationStepHardness::Hardened
338 );
339 assert_eq!(
340 extract_step_hardness("..9h"),
341 ExplorationStepHardness::Hardened
342 );
343 assert_eq!(extract_step_hardness("9"), ExplorationStepHardness::Normal);
344 assert_eq!(
345 extract_step_hardness("*a"),
346 ExplorationStepHardness::HardenedAndNormal
347 );
348 }
349
350 #[test]
351 fn translate_wildcard_step_string_to_exploration_step_works_01() {
352 let result = translate_wildcard_step_string_to_exploration_step("*h".to_string(), 10);
353 let expected = ExplorationStep::new(0, 10, ExplorationStepHardness::Hardened);
354 assert_eq!(result, expected);
355
356 let result = translate_wildcard_step_string_to_exploration_step("*'".to_string(), 10);
357 let expected = ExplorationStep::new(0, 10, ExplorationStepHardness::Hardened);
358 assert_eq!(result, expected);
359
360 let result = translate_wildcard_step_string_to_exploration_step("*a".to_string(), 10);
361 let expected = ExplorationStep::new(0, 10, ExplorationStepHardness::HardenedAndNormal);
362 assert_eq!(result, expected);
363
364 let result = translate_wildcard_step_string_to_exploration_step("*".to_string(), 10);
365 let expected = ExplorationStep::new(0, 10, ExplorationStepHardness::Normal);
366 assert_eq!(result, expected);
367 }
368
369 #[test]
370 fn translate_range_step_string_to_exploration_step_works_01() {
371 let result = translate_range_step_string_to_exploration_step("0".to_string()).unwrap();
372 let expected = ExplorationStep::new(0, 0, ExplorationStepHardness::Normal);
373 assert_eq!(result, expected);
374
375 let result = translate_range_step_string_to_exploration_step("9..78h".to_string()).unwrap();
376 let expected = ExplorationStep::new(9, 78, ExplorationStepHardness::Hardened);
377 assert_eq!(result, expected);
378
379 let result =
380 translate_range_step_string_to_exploration_step("100..120a".to_string()).unwrap();
381 let expected = ExplorationStep::new(100, 120, ExplorationStepHardness::HardenedAndNormal);
382 assert_eq!(result, expected);
383
384 let result = translate_range_step_string_to_exploration_step("..10".to_string()).unwrap();
385 let expected = ExplorationStep::new(0, 10, ExplorationStepHardness::Normal);
386 assert_eq!(result, expected);
387
388 let result = translate_range_step_string_to_exploration_step("9..7".to_string());
389 assert!(result.is_err());
390 }
391
392 #[test]
393 fn new_works_01() {
394 let exploration_str = "0/..8/*h/6..9a/*'/40a";
395 let result = ExplorationPath::new(None, exploration_str, 5, false).unwrap();
396 let expected = ExplorationPath {
397 base_paths: vec![DerivationPath::from_str("m").unwrap()],
398 explore: vec![
399 ExplorationStep::new(0, 0, ExplorationStepHardness::Normal),
400 ExplorationStep::new(0, 8, ExplorationStepHardness::Normal),
401 ExplorationStep::new(0, 5, ExplorationStepHardness::Hardened),
402 ExplorationStep::new(6, 9, ExplorationStepHardness::HardenedAndNormal),
403 ExplorationStep::new(0, 5, ExplorationStepHardness::Hardened),
404 ExplorationStep::new(40, 40, ExplorationStepHardness::HardenedAndNormal),
405 ],
406 depth: 5,
407 sweep: false,
408 };
409 assert_eq!(expected, result);
410 }
411
412 #[test]
413 fn new_works_02() {
414 let exploration_str = "..9a";
415 let result = ExplorationPath::new(None, exploration_str, 5, false).unwrap();
416 let expected = ExplorationPath {
417 base_paths: vec![DerivationPath::from_str("m").unwrap()],
418 explore: vec![ExplorationStep::new(
419 0,
420 9,
421 ExplorationStepHardness::HardenedAndNormal,
422 )],
423 depth: 5,
424 sweep: false,
425 };
426 assert_eq!(result, expected);
427 }
428
429 #[test]
430 fn new_works_03() {
431 let exploration_str = "0u/..8/*h/6..9a/*'/40a";
432 let result = ExplorationPath::new(None, exploration_str, 5, false);
433 assert!(result.is_err())
434 }
435
436 #[test]
437 fn new_works_04() {
438 let exploration_str = "./.8";
439 let result = ExplorationPath::new(None, exploration_str, 5, false);
440 assert!(result.is_err())
441 }
442
443 #[test]
444 fn new_works_05() {
445 let exploration_str = "";
446 let result = ExplorationPath::new(None, exploration_str, 5, false);
447 assert!(result.is_err());
448 }
449
450 #[test]
451 fn num_of_paths_works_01() {
452 let exploration_path = ExplorationPath::new(None, "..8", 5, false).unwrap();
453 assert_eq!(exploration_path.num_of_paths(), 9);
454
455 let exploration_path = ExplorationPath::new(None, "4..8h", 5, false).unwrap();
456 assert_eq!(exploration_path.num_of_paths(), 5);
457
458 let exploration_path = ExplorationPath::new(None, "8'", 5, false).unwrap();
459 assert_eq!(exploration_path.num_of_paths(), 1);
460
461 let exploration_path = ExplorationPath::new(None, "*a", 5, false).unwrap();
462 assert_eq!(exploration_path.num_of_paths(), 12);
463
464 let exploration_path = ExplorationPath::new(None, "..8/*a", 5, false).unwrap();
465 assert_eq!(exploration_path.num_of_paths(), 108);
466
467 let exploration_path = ExplorationPath::new(None, "3..9h/*'/9a/*/*h", 5, false).unwrap();
468 assert_eq!(exploration_path.num_of_paths(), 3024);
469
470 let exploration_path = ExplorationPath::new(None, "/8/*a/..90'/0", 5, false).unwrap();
471 assert_eq!(exploration_path.num_of_paths(), 1092);
472 }
473
474 #[test]
475 fn num_of_paths_sweep_from_root_works_01() {
476 let exploration_path = ExplorationPath::new(None, "*a/..2h/4", 1, false).unwrap();
477 assert_eq!(exploration_path.num_of_paths_sweep(), 29);
478 }
479
480 #[test]
481 fn num_of_paths_sweep_from_root_works_02() {
482 let exploration_path = ExplorationPath::new(None, "*a/..2h/4", 3, false).unwrap();
483 assert_eq!(exploration_path.num_of_paths_sweep(), 57);
484 }
485}