feat: parse rule and mapping files to determine file extensions to load

src/file/mod.rs

@@ -1,3 +1,4 @@
+use std::collections::HashSet;
 use std::fs;
 use std::path::{Path, PathBuf};
 
@@ -22,7 +23,7 @@ pub struct Documents<'a> {
     iterator: Box<dyn Iterator<Item = crate::Result<Document>> + 'a>,
 }
 
-#[derive(Clone, Debug, PartialEq, Deserialize, Serialize)]
+#[derive(Clone, Debug, PartialEq, Deserialize, Serialize, Hash, Eq)]
 #[serde(rename_all = "snake_case")]
 pub enum Kind {
     Evtx,
@@ -31,6 +32,17 @@ pub enum Kind {
     Unknown,
 }
 
+impl Kind {
+    pub fn extensions(&self) -> Option<Vec<String>> {
+        match self {
+            Kind::Evtx => Some(vec!["evtx".to_string()]),
+            Kind::Json => Some(vec!["json".to_string()]),
+            Kind::Xml => Some(vec!["xml".to_string()]),
+            Kind::Unknown => None,
+        }
+    }
+}
+
 impl<'a> Iterator for Documents<'a> {
     type Item = crate::Result<Document>;
 
@@ -86,7 +98,7 @@ impl Reader {
                             })
                         } else {
                             anyhow::bail!(
-                                "file type is not currently supported - {}",
+                                "file type is not currently supported - {}, use --skip-errors to continue",
                                 file.display()
                             )
                         }
@@ -161,7 +173,7 @@ impl Reader {
 
 pub fn get_files(
     path: &PathBuf,
-    extension: &Option<String>,
+    extensions: &Option<HashSet<String>>,
     skip_errors: bool,
 ) -> crate::Result<Vec<PathBuf>> {
     let mut files: Vec<PathBuf> = vec![];
@@ -201,11 +213,11 @@ pub fn get_files(
                         }
                     }
                 };
-                files.extend(get_files(&dir.path(), extension, skip_errors)?);
+                files.extend(get_files(&dir.path(), extensions, skip_errors)?);
             }
-        } else if let Some(extension) = extension {
+        } else if let Some(e) = extensions {
             if let Some(ext) = path.extension() {
-                if ext == extension.as_str() {
+                if e.contains(&ext.to_string_lossy().into_owned()) {
                     files.push(path.to_path_buf());
                 }
             }

src/hunt.rs

@@ -359,8 +359,8 @@ impl Hunt {
 }
 
 pub struct HunterInner {
-    hunts: Vec<Hunt>,
-    rules: BTreeMap<Uuid, Rule>,
+    pub hunts: Vec<Hunt>,
+    pub rules: BTreeMap<Uuid, Rule>,
 
     load_unknown: bool,
     local: bool,
@@ -371,7 +371,7 @@ pub struct HunterInner {
 }
 
 pub struct Hunter {
-    inner: HunterInner,
+    pub inner: HunterInner,
 }
 
 impl Hunter {
@@ -607,6 +607,20 @@ impl Hunter {
     pub fn rules(&self) -> &BTreeMap<Uuid, Rule> {
         &self.inner.rules
     }
+    pub fn extensions(&self) -> HashSet<String> {
+        let mut extensions = HashSet::new();
+        for rule in &self.inner.rules {
+            if let Some(e) = FileKind::extensions(&rule.1.types()) {
+                extensions.extend(e.iter().cloned());
+            }
+        }
+        for hunt in &self.inner.hunts {
+            if let Some(e) = FileKind::extensions(&hunt.file) {
+                extensions.extend(e.iter().cloned());
+            }
+        }
+        extensions
+    }
 
     fn skip(&self, timestamp: NaiveDateTime) -> crate::Result<bool> {
         if self.inner.from.is_some() || self.inner.to.is_some() {

src/lib.rs

@@ -3,7 +3,7 @@ extern crate anyhow;
 
 pub(crate) use anyhow::Result;
 
-pub use file::{evtx, get_files, Reader};
+pub use file::{evtx, get_files, Kind as FileKind, Reader};
 pub use hunt::{Hunter, HunterBuilder};
 pub use rule::{
     lint, load, sigma, Filter, Kind as RuleKind, Level as RuleLevel, Status as RuleStatus,

src/main.rs

@@ -54,8 +54,8 @@ enum Command {
         #[structopt(group = "format", long = "csv", requires("output"))]
         csv: bool,
         /// Only hunt through files with the provided extension.
-        #[structopt(long = "extension")]
-        extension: Option<String>,
+        #[structopt(long = "extension", number_of_values = 1)]
+        extension: Option<Vec<String>>,
         /// The timestamp to hunt from. Drops any documents older than the value provided.
         #[structopt(long = "from")]
         from: Option<NaiveDateTime>,
@@ -129,8 +129,8 @@ enum Command {
         regexp: Option<Vec<String>>,
 
         /// Only search through files with the provided extension.
-        #[structopt(long = "extension")]
-        extension: Option<String>,
+        #[structopt(long = "extension", number_of_values = 1)]
+        extension: Option<Vec<String>>,
         /// The timestamp to search from. Drops any documents older than the value provided.
         #[structopt(long = "from")]
         from: Option<NaiveDateTime>,
@@ -267,14 +267,6 @@ fn run() -> Result<()> {
             };
             let sigma = sigma.unwrap_or_default();
 
-            cs_eprintln!(
-                "[+] Loading event logs from: {}",
-                path.iter()
-                    .map(|p| p.display().to_string())
-                    .collect::<Vec<_>>()
-                    .join(", ")
-            );
-
             cs_eprintln!(
                 "[+] Loading detection rules from: {}",
                 rules
@@ -364,10 +356,54 @@ fn run() -> Result<()> {
                 hunter = hunter.to(to);
             }
             let hunter = hunter.build()?;
+
+            /* if no user-defined extensions are specified, then we parse rules and
+            mappings to build a list of file extensions that should be loaded */
+            let mut scratch = HashSet::new();
+            let message;
+            let exts = if load_unknown {
+                message = "*".to_string();
+                None
+            } else {
+                scratch.extend(hunter.extensions());
+                if scratch.is_empty() {
+                    return Err(anyhow::anyhow!(
+                        "No valid file extensions for the 'kind' specified in the mapping or rules files"
+                    ));
+                }
+                if let Some(e) = extension {
+                    // User has provided specific extensions
+                    scratch = scratch
+                        .intersection(&HashSet::from_iter(e.iter().cloned()))
+                        .cloned()
+                        .collect();
+                    if scratch.is_empty() {
+                        return Err(anyhow::anyhow!(
+                        "The specified file extension is not supported. Use --load-unknown to force loading",
+                    ));
+                    }
+                };
+                message = scratch
+                    .iter()
+                    .map(|x| format!(".{}", x))
+                    .collect::<Vec<_>>()
+                    .join(", ");
+                Some(scratch)
+            };
+
+            cs_eprintln!(
+                "[+] Loading event logs from: {} (extensions: {})",
+                path.iter()
+                    .map(|p| p.display().to_string())
+                    .collect::<Vec<_>>()
+                    .join(", "),
+                message
+            );
+
             let mut files = vec![];
             let mut size = ByteSize::mb(0);
             for path in &path {
-                let res = get_files(path, &extension, skip_errors)?;
+                let res = get_files(path, &exts, skip_errors)?;
                 for i in &res {
                     size += i.metadata()?.len();
                 }
@@ -509,16 +545,46 @@ fn run() -> Result<()> {
                     std::env::current_dir().expect("could not get current working directory"),
                 );
             }
+            let types = if let Some(e) = &extension {
+                Some(HashSet::from_iter(e.clone()))
+            } else {
+                None
+            };
+
             let mut files = vec![];
             let mut size = ByteSize::mb(0);
             for path in &paths {
-                let res = get_files(path, &extension, skip_errors)?;
+                let res = get_files(path, &types, skip_errors)?;
                 for i in &res {
                     size += i.metadata()?.len();
                 }
                 files.extend(res);
             }
 
+            if let Some(ext) = &extension {
+                cs_eprintln!(
+                    "[+] Loading event logs from: {} (extensions: {})",
+                    paths
+                        .iter()
+                        .map(|p| p.display().to_string())
+                        .collect::<Vec<_>>()
+                        .join(", "),
+                    ext.iter()
+                        .map(|x| format!(".{}", x))
+                        .collect::<Vec<_>>()
+                        .join(", ")
+                )
+            } else {
+                cs_eprintln!(
+                    "[+] Loading event logs from: {}",
+                    paths
+                        .iter()
+                        .map(|p| p.display().to_string())
+                        .collect::<Vec<_>>()
+                        .join(", "),
+                )
+            };
+
             if files.is_empty() {
                 return Err(anyhow::anyhow!(
                     "No event logs were found in the provided paths",

src/rule/mod.rs

@@ -1,3 +1,4 @@
+use crate::file::Kind as FileKind;
 use std::collections::HashSet;
 use std::fmt;
 use std::path::Path;
@@ -50,6 +51,14 @@ impl Rule {
         }
     }
 
+    #[inline]
+    pub fn types(&self) -> &FileKind {
+        match self {
+            Self::Chainsaw(c) => &c.kind,
+            Self::Sigma(_) => &FileKind::Unknown,
+        }
+    }
+
     #[inline]
     pub fn name(&self) -> &String {
         match self {