added add-file-extensions option #586

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,

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,22 +186,28 @@ 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() {
  return;
  }
  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);
  }
  }