diff options
author | Omer Ben-Amram <omerbenamram@gmail.com> | 2019-06-06 10:37:35 +0300 |
---|---|---|
committer | Omer Ben-Amram <omerbenamram@gmail.com> | 2019-06-06 10:37:35 +0300 |
commit | 3d1861c50cb4a742e6b163d5a9fa271d9aff4b34 (patch) | |
tree | 0ea803d656a03efe24b24613a4f4fc253f94fe71 | |
parent | c2e59b9c88f28ff79dca7f166c1d911e4600976a (diff) |
implement ability to select entry ranges
-rw-r--r-- | src/bin/mft_dump.rs | 114 |
1 files changed, 111 insertions, 3 deletions
diff --git a/src/bin/mft_dump.rs b/src/bin/mft_dump.rs index e699f36..628bd6c 100644 --- a/src/bin/mft_dump.rs +++ b/src/bin/mft_dump.rs @@ -12,13 +12,14 @@ use mft::csv::FlatMftEntryWithName; use snafu::ErrorCompat; use std::fs::File; -use std::io::Read; use std::io::Write; use std::path::{Path, PathBuf}; use std::process::exit; use mft::entry::ZERO_HEADER; use std::fmt::Write as FmtWrite; +use std::ops::RangeInclusive; +use std::str::FromStr; use std::{fs, io, path}; /// Simple error macro for use inside of internal errors in `MftDump` @@ -46,6 +47,89 @@ impl OutputFormat { } } +struct Ranges(Vec<RangeInclusive<usize>>); + +impl Ranges { + pub fn chain(&self) -> impl Iterator<Item = usize> + '_ { + self.0.iter().cloned().flatten() + } +} + +impl FromStr for Ranges { + type Err = StdErr; + + fn from_str(s: &str) -> Result<Self, Self::Err> { + let mut ranges = vec![]; + for x in s.split(',') { + // range + if x.contains('-') { + let range: Vec<&str> = x.split('-').collect(); + if range.len() != 2 { + return err!( + "Failed to parse ranges: Range should contain exactly one `-`, found {}", + x + ); + } + + ranges.push(range[0].parse()?..=range[1].parse()?); + } else { + let n = x.parse()?; + ranges.push(n..=n); + } + } + + Ok(Ranges(ranges)) + } +} + +#[cfg(test)] +mod tests { + use super::Ranges; + use std::str::FromStr; + + #[test] + fn it_works_with_single_number() { + let ranges = Ranges::from_str("1").unwrap(); + assert_eq!(ranges.0, vec![1..=1]); + } + + #[test] + fn it_works_with_a_range() { + let ranges = Ranges::from_str("1-5").unwrap(); + assert_eq!(ranges.0, vec![1..=5]); + } + + #[test] + fn it_works_with_a_range_and_a_number() { + let ranges = Ranges::from_str("1-5,8").unwrap(); + assert_eq!(ranges.0, vec![1..=5, 8..=8]); + } + + #[test] + fn it_works_with_a_number_and_a_range() { + let ranges = Ranges::from_str("1-5,8").unwrap(); + assert_eq!(ranges.0, vec![1..=5, 8..=8]); + } + + #[test] + fn it_works_with_more_than_2_number_and_a_range() { + let ranges = Ranges::from_str("1-5,8,10-19").unwrap(); + assert_eq!(ranges.0, vec![1..=5, 8..=8, 10..=19]); + } + + #[test] + fn it_errors_on_a_random_string() { + let ranges = Ranges::from_str("hello"); + assert!(ranges.is_err()) + } + + #[test] + fn it_errors_on_a_range_with_too_many_dashes() { + let ranges = Ranges::from_str("1-5-8"); + assert!(ranges.is_err()) + } +} + struct MftDump { filepath: PathBuf, // We use an option here to be able to move the output out of mftdump from a mutable reference. @@ -53,6 +137,7 @@ struct MftDump { data_streams_output: Option<PathBuf>, verbosity_level: Option<Level>, output_format: OutputFormat, + ranges: Option<Ranges>, backtraces: bool, } @@ -98,12 +183,18 @@ impl MftDump { } }; + let ranges = match matches.value_of("entry-range") { + Some(range) => Some(Ranges::from_str(range)?), + None => None, + }; + Ok(MftDump { filepath: PathBuf::from(matches.value_of("INPUT").expect("Required argument")), output, data_streams_output, verbosity_level, output_format, + ranges, backtraces, }) } @@ -203,8 +294,18 @@ impl MftDump { }; let number_of_entries = parser.get_entry_count(); - for i in 0..number_of_entries { - let entry = parser.get_entry(i); + + // Move ranges out of self here to avoid immutably locking self during + // the `for i in entries` loop. + let take_ranges = self.ranges.take(); + + let entries = match take_ranges { + Some(ref ranges) => Box::new(ranges.chain()), + None => Box::new(0..number_of_entries as usize) as Box<dyn Iterator<Item = usize>>, + }; + + for i in entries { + let entry = parser.get_entry(i as u64); let entry = match entry { Ok(entry) => match &entry.header.signature { @@ -381,6 +482,13 @@ fn main() { .help("Output format."), ) .arg( + Arg::with_name("entry-range") + .long("--ranges") + .short("-r") + .takes_value(true) + .help(indoc!("Dumps only the given entry range(s), for example, `1-15,30` will dump entries 1-15, and 30")), + ) + .arg( Arg::with_name("output-target") .long("--output") .short("-f") |