Merge pull request #38 from dfir-dd/37-error-in-hivescan-unexpectedeof

37 error in hivescan unexpectedeof

.gitignore

@@ -1,3 +1,6 @@
 target
 !Cargo.lock
 .vscode
+.idea
+tmp
+temp

Cargo.toml

@@ -171,7 +171,8 @@ strum = { version = "0", features = ["derive"], optional=true }
 strum_macros = {version="0", optional=true}
 
 # nt-hive2
-nt_hive2 = {version="4.1.0", optional=true}
+nt_hive2 = {version="4.2.1", optional=true}
+#nt_hive2 = {path="../nt-hive2", optional=true}
 
 # lnk2bodyfile
 lnk = {version="0.5.1", optional=true}

src/bin/hivescan/hivescanapplication.rs

@@ -4,21 +4,26 @@ use anyhow::Result;
 use dfir_toolkit::common::bodyfile::Bodyfile3Line;
 use indicatif::{ProgressBar, ProgressStyle};
 use nt_hive2::*;
-use std::fs::File;
-use std::time::Duration;
+use std::{io::Read, io::Seek};
 
 use crate::regtreebuilder::RegTreeBuilder;
 
-pub(crate) struct HiveScanApplication {
+pub(crate) struct HiveScanApplication<RS>
+where
+    RS: Read + Seek,
+{
     #[allow(dead_code)]
     cli: Cli,
 
     root_offset: Offset,
-    hive: Option<Hive<File, CleanHive>>,
+    hive: Option<Hive<RS, CleanHive>>,
 }
 
-impl HiveScanApplication {
-    pub fn new(cli: Cli, hive: Hive<File, CleanHive>) -> Self {
+impl<RS> HiveScanApplication<RS>
+where
+    RS: Read + Seek,
+{
+    pub fn new(cli: Cli, hive: Hive<RS, CleanHive>) -> Self {
         Self {
             cli,
             root_offset: hive.root_cell_offset(),
@@ -34,7 +39,6 @@ impl HiveScanApplication {
             .unwrap();
         let bar = ProgressBar::new(self.hive.as_ref().unwrap().data_size().into());
         bar.set_style(progress_style);
-        bar.enable_steady_tick(Duration::from_millis(100));
         bar.set_message("scanning cells");
 
         let builder = RegTreeBuilder::from_hive(self.hive.take().unwrap(), |p| bar.set_position(p));

src/bin/hivescan/main.rs

@@ -2,7 +2,7 @@ use anyhow::{bail, Result};
 use cli::Cli;
 use dfir_toolkit::common::FancyParser;
 use nt_hive2::{ContainsHive, Hive, HiveParseMode};
-use std::fs::File;
+use std::{fs::File, io::BufReader};
 
 mod hivescanapplication;
 mod regtreebuilder;
@@ -15,6 +15,7 @@ fn main() -> Result<()> {
 
     match File::open(&cli.hive_file) {
         Ok(data) => {
+            let data = BufReader::new(data);
             let hive = Hive::new(data, HiveParseMode::NormalWithBaseBlock).unwrap();
 
             let clean_hive = match cli.logfiles.len() {

src/bin/hivescan/regtreebuilder.rs

@@ -1,7 +1,5 @@
 use std::{
-    cell::RefCell,
-    collections::{hash_map, HashMap},
-    rc::Rc,
+    cell::RefCell, cmp::max, collections::{hash_map, HashMap}, rc::Rc
 };
 
 use binread::BinReaderExt;
@@ -48,25 +46,28 @@ impl RegTreeBuilder {
         B: BinReaderExt,
         C: Fn(u64),
     {
-        let iterator = hive.into_cell_iterator(progress_callback);
         let mut me = Self {
             subtrees: HashMap::new(),
             entries: HashMap::new(),
             missing_parents: HashMap::new(),
         };
 
         let mut last_offset = Offset(0);
-        for cell in iterator {
+        let mut last_cell_end = 0;
+        for cell in hive.hivebins().flat_map(|hivebin| hivebin.cells()) {
             let my_offset = *cell.offset();
+            let my_size = cell.header().size();
             let is_deleted = cell.header().is_deleted();
             assert_ne!(last_offset, my_offset);
             log::trace!("found new cell at offset 0x{:x}", my_offset.0);
 
             if let Ok(nk) = TryInto::<KeyNode>::try_into(cell) {
-                me.insert_nk(my_offset, nk, is_deleted)
-            };
-
+                me.insert_nk(my_offset, nk, is_deleted);
+            }
             last_offset = my_offset;
+
+            last_cell_end = max(last_cell_end, u64::from(last_offset.0) + u64::try_from(my_size).unwrap());
+            progress_callback(last_cell_end);
         }
         me
     }
@@ -78,7 +79,7 @@ impl RegTreeBuilder {
     }
 
     fn insert_nk(&mut self, nk_offset: Offset, nk: KeyNode, is_deleted: bool) {
-        assert!(!self.subtrees.contains_key(&nk_offset));
+        assert!(!self.subtrees.contains_key(&nk_offset), "KeyNode at offset 0x{:08x} is already in the set of subtrees", nk_offset.0);
         assert!(!self.entries.contains_key(&nk_offset));
 
         let parent_offset = nk.parent;

tests/hivescan/mod.rs

@@ -0,0 +1,2 @@
+mod testhive;
+mod new_dirty_hive1;

tests/hivescan/new_dirty_hive1.rs

@@ -0,0 +1,15 @@
+use std::path::PathBuf;
+
+use assert_cmd::Command;
+
+#[test]
+fn scan_new_dirty_hive1() {
+    let mut cmd = Command::cargo_bin("hivescan").unwrap();
+    let mut data_path = PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap());
+    data_path.push("tests");
+    data_path.push("data");
+    data_path.push("hivescan");
+    data_path.push("NewDirtyHive1");
+    data_path.push("NewDirtyHive");
+    cmd.arg(data_path).assert().success();
+}

tests/hivescan/testhive.rs

@@ -0,0 +1,14 @@
+use std::path::PathBuf;
+
+use assert_cmd::Command;
+
+#[test]
+fn scan_testhive() {
+    let mut cmd = Command::cargo_bin("hivescan").unwrap();
+    let mut data_path = PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap());
+    data_path.push("tests");
+    data_path.push("data");
+    data_path.push("hivescan");
+    data_path.push("testhive");
+    cmd.arg(data_path).assert().success();
+}

tests/mod.rs

@@ -1,3 +1,4 @@
 mod mactime2;
 mod ts2date;
-mod lnk2bodyfile;
+mod lnk2bodyfile;
+mod hivescan;