Added jmes query functionality

2 files changed, 66 insertions, 15 deletions

diff --git a/Cargo.toml b/Cargo.toml
index e2d4717..28f3600 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -18,6 +18,9 @@ serde_derive = "0.9"
 serde_json = "0.9"
 seek_bufread = "~1.2"
 
+[dependencies.jmespath]
+version = "^0.1.1"
+
 [dependencies.r-winstructs]
 version = "*"
 branch = "master"
diff --git a/src/main.rs b/src/main.rs
index c4c68d5..e112f8c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,10 +2,12 @@
 extern crate rustymft;
 extern crate rwinstructs;
 extern crate serde_json;
+extern crate jmespath;
 extern crate serde;
 extern crate clap;
-use clap::{App, Arg};
+use clap::{App, Arg, ArgMatches};
 use rustymft::mft::{MftHandler};
+use jmespath::{Expression};
 use rwinstructs::reference;
 use rwinstructs::serialize;
 use serde::Serializer;
@@ -13,22 +15,34 @@ use serde::ser::SerializeSeq;
 use std::fs;
 use std::io;
 
-fn process_directory<S>(directory: &str, serializer: S) where S: serde::Serializer {
-    let mut seq = serializer.serialize_seq(None).unwrap();
+fn process_directory(directory: &str, options: ArgMatches) {
     for entry in fs::read_dir(directory).unwrap() {
         let entry = entry.unwrap();
         let path = entry.path();
         if path.is_file() {
             let path_string = path.into_os_string().into_string().unwrap();
             if path_string.ends_with("mft"){
-                process_file(&path_string,&mut seq);
+                process_file(&path_string,options.clone());
             }
         }
     }
-    seq.end().unwrap();
 }
 
-fn process_file<S: serde::ser::SerializeSeq>(filename: &str, serializer: &mut S) -> bool {
+fn process_file(filename: &str,options: ArgMatches) -> bool {
+    // JMES Expression if needed
+    let mut expr: Option<Expression> = None;
+    if options.is_present("query") {
+        expr = Some(jmespath::compile(
+            options.value_of("query").unwrap()
+        ).unwrap());
+    }
+
+    // Expression bool flag
+    let mut expr_as_bool = false;
+    if options.is_present("bool_expr"){
+        expr_as_bool = true;
+    }
+
     let mut mft_handler = match MftHandler::new(filename) {
         Ok(mft_handler) => mft_handler,
         Err(error) => {
@@ -39,13 +53,39 @@ fn process_file<S: serde::ser::SerializeSeq>(filename: &str, serializer: &mut S)
 
     for i in 0 .. mft_handler.get_entry_count() {
         let mft_entry = match mft_handler.entry(i) {
-            Ok(mft_entry) => mft_entry,
+            Ok(mft_entry) => {
+                let json_str = serde_json::to_string(&mft_entry).unwrap();
+
+                match expr {
+                    Some(ref j_expr) => {
+                        let data = jmespath::Variable::from_json(&json_str).unwrap();
+                        let result = j_expr.search(data).unwrap();
+                        if expr_as_bool {
+                            match result.as_boolean() {
+                                Some(bool_value) => {
+                                    match bool_value {
+                                        true => println!("{}",json_str),
+                                        false => {}
+                                    }
+                                },
+                                None => {
+                                    panic!("Query expression is not a bool expression!");
+                                }
+                            }
+                        } else {
+                            println!("{}",result)
+                        }
+                    },
+                    None => {
+                        println!("{}",json_str);
+                    }
+                }
+            },
             Err(error) => {
                 warn!("Could not parse file: {} [error: {}]", filename, error);
                 continue;
             }
         };
-        serializer.serialize_element(&mft_entry).unwrap();
     }
 
     return true;
@@ -64,11 +104,25 @@ fn main() {
         .required(true)
         .takes_value(true);
 
+    let jmes_arg = Arg::with_name("query")
+        .short("q")
+        .long("query")
+        .value_name("QUERY")
+        .help("JMES Query")
+        .takes_value(true);
+
+    let bool_arg = Arg::with_name("bool_expr")
+        .short("b")
+        .long("bool_expr")
+        .help("JMES Query as bool only. (Prints whole record if true.)");
+
     let options = App::new("RustyMft")
         .version("0.0.0")
         .author("Matthew Seyer <https://github.com/forensicmatt/RustyMft>")
         .about("Parse $MFT.")
         .arg(source_arg)
+        .arg(jmes_arg)
+        .arg(bool_arg)
         .get_matches();
 
     // Set Reference Display Options
@@ -77,15 +131,9 @@ fn main() {
 
     let source = options.value_of("source").unwrap();
 
-    let mut serializer = serde_json::Serializer::pretty(
-        io::stdout()
-    );
-
     if is_directory(source) {
         panic!("Directory source is not implemented yet.");
     } else {
-        let mut seq = serializer.serialize_seq(None).unwrap();
-        process_file(source,&mut seq);
-        seq.end().unwrap();
+        process_file(source,options.clone());
     }
 }