From 21dbe2c97a1a775a158e78ff0251eb3b48d0cbd8 Mon Sep 17 00:00:00 2001
From: DustInDark <nextsasasa@gmail.com>
Date: Mon, 20 Jun 2022 19:53:44 +0900
Subject: [PATCH] added add-file-extensions option #586

---
 src/detections/configs.rs | 14 ++++++++++++
 src/main.rs               | 47 ++++++++++++++++++++++-----------------
 2 files changed, 41 insertions(+), 20 deletions(-)

diff --git a/src/detections/configs.rs b/src/detections/configs.rs
index 3da13610..18bb3489 100644
--- a/src/detections/configs.rs
+++ b/src/detections/configs.rs
@@ -30,6 +30,8 @@ lazy_static! {
     pub static ref IDS_REGEX: Regex =
         Regex::new(r"^[0-9a-z]{8}-[0-9a-z]{4}-[0-9a-z]{4}-[0-9a-z]{4}-[0-9a-z]{12}$").unwrap();
     pub static ref TERM_SIZE: Option<(Width, Height)> = terminal_size();
+    pub static ref TARGET_EXTENSIONS: HashSet<String> =
+        get_target_extensions(CONFIG.read().unwrap().args.add_file_extentions.as_ref());
 }
 
 pub struct ConfigReader<'a> {
@@ -205,6 +207,10 @@ pub struct Config {
     /// Print the list of contributors
     #[clap(long)]
     pub contributors: bool,
+
+    /// Specify target file extension expclude evtx (ex: evtx_data)
+    #[clap(long = "add-file-extensions", multiple_values = true)]
+    pub add_file_extentions: Option<Vec<String>>,
 }
 
 impl ConfigReader<'_> {
@@ -453,6 +459,14 @@ pub fn load_pivot_keywords(path: &str) {
     });
 }
 
+/// --add-file-extensionsで追加された拡張子から、調査対象ファイルの拡張子セットを返す関数
+pub fn get_target_extensions(arg: Option<&Vec<String>>) -> HashSet<String> {
+    let mut target_file_extensions: HashSet<String> =
+        arg.unwrap_or(&Vec::new()).iter().cloned().collect();
+    target_file_extensions.insert(String::from("evtx"));
+    target_file_extensions
+}
+
 #[derive(Debug, Clone)]
 pub struct EventInfo {
     pub evttitle: String,
diff --git a/src/main.rs b/src/main.rs
index fcd72e05..543f1b6a 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -11,7 +11,7 @@ use chrono::{DateTime, Datelike, Local, TimeZone};
 use evtx::{EvtxParser, ParserSettings};
 use git2::Repository;
 use hashbrown::{HashMap, HashSet};
-use hayabusa::detections::configs::{load_pivot_keywords, TargetEventTime};
+use hayabusa::detections::configs::{load_pivot_keywords, TargetEventTime, TARGET_EXTENSIONS};
 use hayabusa::detections::detection::{self, EvtxRecordInfo};
 use hayabusa::detections::pivot::PivotKeyword;
 use hayabusa::detections::pivot::PIVOT_KEYWORD;
@@ -186,6 +186,7 @@ impl App {
             .ok();
             println!();
         }
+
         if configs::CONFIG.read().unwrap().args.live_analysis {
             let live_analysis_list = self.collect_liveanalysis_files();
             if live_analysis_list.is_none() {
@@ -193,15 +194,20 @@ impl App {
             }
             self.analysis_files(live_analysis_list.unwrap(), &time_filter);
         } else if let Some(filepath) = &configs::CONFIG.read().unwrap().args.filepath {
-            if filepath.extension().unwrap_or_else(|| OsStr::new(".")) != "evtx"
-                || filepath
-                    .as_path()
-                    .file_stem()
+            if TARGET_EXTENSIONS.contains(
+                filepath
+                    .extension()
                     .unwrap_or_else(|| OsStr::new("."))
                     .to_str()
-                    .unwrap()
-                    .trim()
-                    .starts_with('.')
+                    .unwrap(),
+            ) || filepath
+                .as_path()
+                .file_stem()
+                .unwrap_or_else(|| OsStr::new("."))
+                .to_str()
+                .unwrap()
+                .trim()
+                .starts_with('.')
             {
                 AlertMessage::alert(
                     "--filepath only accepts .evtx files. Hidden files are ignored.",
@@ -397,18 +403,19 @@ impl App {
                     ret.extend(subdir_ret);
                     Option::Some(())
                 });
-            } else {
-                let path_str = path.to_str().unwrap_or("");
-                if path_str.ends_with(".evtx")
-                    && !Path::new(path_str)
-                        .file_stem()
-                        .unwrap_or_else(|| OsStr::new("."))
-                        .to_str()
-                        .unwrap()
-                        .starts_with('.')
-                {
-                    ret.push(path);
-                }
+            } else if TARGET_EXTENSIONS.contains(
+                path.extension()
+                    .unwrap_or_else(|| OsStr::new(""))
+                    .to_str()
+                    .unwrap(),
+            ) && !path
+                .file_stem()
+                .unwrap_or_else(|| OsStr::new("."))
+                .to_str()
+                .unwrap()
+                .starts_with('.')
+            {
+                ret.push(path);
             }
         }