Code for Episode 9

Author: h33p

scanflow/src/main.rs

@@ -7,6 +7,7 @@ use memflow_win32::{Error, Result};
 
 use simplelog::{Config, LevelFilter, TermLogger, TerminalMode};
 
+use std::collections::BTreeSet;
 use std::convert::TryInto;
 
 mod value_scanner;
@@ -46,19 +47,12 @@ fn main() -> Result<()> {
     //let replace_str = b"Hello world from memflow!";
 
     let mut value_scanner = ValueScanner::default();
-
     let mut typename: Option<String> = None;
     let mut buf_len = 0;
 
     let mut disasm = Disasm::default();
-    disasm.collect_globals(&mut process)?;
-
     let mut pointer_map = PointerMap::default();
 
-    for (&instr, &global) in disasm.map().iter().filter(|(_, &o)| o == 0xa9770.into()) {
-        println!("{:x} -> {:x}", instr, global);
-    }
-
     while let Ok(line) = get_line() {
         let line = line.trim();
 
@@ -87,9 +81,14 @@ fn main() -> Result<()> {
                     process.proc_info.proc_arch.size_addr(),
                 )?;
             }
+            "globals" | "g" => {
+                disasm.reset();
+                disasm.collect_globals(&mut process)?;
+                println!("Global variable references found: {:x}", disasm.map().len());
+            }
             "offset_scan" | "os" => {
-                if let (Some(lrange), Some(urange), Some(max_depth), filter_addr) =
-                    scan_fmt_some!(args, "{} {} {} {x}", usize, usize, usize, [hex u64])
+                if let (Some(use_di), Some(lrange), Some(urange), Some(max_depth), filter_addr) =
+                    scan_fmt_some!(args, "{} {} {} {} {x}", String, usize, usize, usize, [hex u64])
                 {
                     if pointer_map.map().is_empty() {
                         pointer_map.create_map(
@@ -98,12 +97,21 @@ fn main() -> Result<()> {
                         )?;
                     }
 
-                    let matches = pointer_map.find_matches(
-                        lrange,
-                        urange,
-                        max_depth,
-                        value_scanner.matches(),
-                    );
+                    let matches = if use_di == "y" {
+                        if disasm.map().is_empty() {
+                            disasm.collect_globals(&mut process)?;
+                        }
+                        let set: BTreeSet<_> = disasm.map().values().copied().collect();
+                        pointer_map.find_matches_addrs(
+                            lrange,
+                            urange,
+                            max_depth,
+                            value_scanner.matches(),
+                            set.into_iter(),
+                        )
+                    } else {
+                        pointer_map.find_matches(lrange, urange, max_depth, value_scanner.matches())
+                    };
 
                     println!("Matches found: {}", matches.len());
 
@@ -124,7 +132,10 @@ fn main() -> Result<()> {
                         println!("{:x}", m);
                     }
                 } else {
-                    println!("usage: os {{lower range}} {{upper range}} {{max depth}}");
+                    println!(
+                        "usage: os {{y/[n]}} {{lower range}} {{upper range}} {{max
+                            depth}} ({{filter}})"
+                    );
                 }
             }
             _ => {

scanflow/src/pointer_map.rs

@@ -98,15 +98,16 @@ impl PointerMap {
         }
     }
 
-    pub fn find_matches(
+    pub fn find_matches_addrs(
         &self,
         lrange: usize,
         urange: usize,
         max_depth: usize,
         search_for: &[Address],
+        addrs: impl Iterator<Item = Address>,
     ) -> Vec<(Address, Vec<(Address, isize)>)> {
         let mut matches = vec![];
-        for &k in self.map.keys() {
+        for k in addrs {
             self.walk_down_range(
                 k,
                 lrange,
@@ -120,6 +121,22 @@ impl PointerMap {
         }
         matches
     }
+
+    pub fn find_matches(
+        &self,
+        lrange: usize,
+        urange: usize,
+        max_depth: usize,
+        search_for: &[Address],
+    ) -> Vec<(Address, Vec<(Address, isize)>)> {
+        self.find_matches_addrs(
+            lrange,
+            urange,
+            max_depth,
+            search_for,
+            self.map.keys().copied(),
+        )
+    }
 }
 
 pub fn signed_diff(a: Address, b: Address) -> isize {