diff --git a/src/afterfact.rs b/src/afterfact.rs index 51c49bb3..33e32510 100644 --- a/src/afterfact.rs +++ b/src/afterfact.rs @@ -16,18 +16,22 @@ pub struct CsvFormat<'a> { } pub fn after_fact() { - let mut target: Box = - if let Some(csv_path) = configs::singleton().args.value_of("csv-timeline") { - match File::create(csv_path) { - Ok(file) => Box::new(file), - Err(err) => { - println!("Failed to open file. {}", err); - process::exit(1); - } + let mut target: Box = if let Some(csv_path) = configs::CONFIG + .read() + .unwrap() + .args + .value_of("csv-timeline") + { + match File::create(csv_path) { + Ok(file) => Box::new(file), + Err(err) => { + println!("Failed to open file. {}", err); + process::exit(1); } - } else { - Box::new(io::stdout()) - }; + } + } else { + Box::new(io::stdout()) + }; if let Err(err) = emit_csv(&mut target) { println!("Failed to write CSV. {}", err); @@ -53,7 +57,7 @@ fn emit_csv(writer: &mut Box) -> Result<(), Box> { } fn format_time(time: &DateTime) -> String { - if configs::singleton().args.is_present("utc") { + if configs::CONFIG.read().unwrap().args.is_present("utc") { format_rfc(time) } else { format_rfc(&time.with_timezone(&Local)) @@ -64,7 +68,7 @@ fn format_rfc(time: &DateTime) -> String where Tz::Offset: std::fmt::Display, { - if configs::singleton().args.is_present("rfc-2822") { + if configs::CONFIG.read().unwrap().args.is_present("rfc-2822") { return time.to_rfc2822(); } else { return time.to_rfc3339(); diff --git a/src/detections/configs.rs b/src/detections/configs.rs index c9b51ddd..896b8aaa 100644 --- a/src/detections/configs.rs +++ b/src/detections/configs.rs @@ -1,32 +1,28 @@ use crate::detections::utils; use clap::{App, AppSettings, Arg, ArgMatches}; +use lazy_static::lazy_static; use std::collections::HashMap; -use std::sync::Once; +use std::sync::RwLock; +lazy_static! { + pub static ref CONFIG: RwLock = RwLock::new(ConfigReader::new()); +} + #[derive(Clone)] -pub struct SingletonReader { +pub struct ConfigReader { pub args: ArgMatches<'static>, pub event_key_alias_config: EventKeyAliasConfig, } -pub fn singleton() -> Box { - static mut SINGLETON: Option> = Option::None; - static ONCE: Once = Once::new(); - - unsafe { - ONCE.call_once(|| { - let singleton = SingletonReader { - args: build_app().get_matches(), - event_key_alias_config: load_eventkey_alias(), - }; - - SINGLETON = Some(Box::new(singleton)); - }); - - return SINGLETON.clone().unwrap(); +impl ConfigReader { + pub fn new() -> Self { + ConfigReader { + args: build_app(), + event_key_alias_config: load_eventkey_alias("config/eventkey_alias.txt"), + } } } -fn build_app() -> clap::App<'static, 'static> { +fn build_app<'a>() -> ArgMatches<'a> { let program = std::env::args() .nth(0) .and_then(|s| { @@ -36,7 +32,11 @@ fn build_app() -> clap::App<'static, 'static> { }) .unwrap(); - App::new(program) + if is_test_mode() { + return ArgMatches::default(); + } + + App::new(&program) .about("Yea! (Yamato Event Analyzer). Aiming to be the world's greatest Windows event log analysis tool!") .version("0.0.1") .author("Author name ") @@ -51,6 +51,17 @@ fn build_app() -> clap::App<'static, 'static> { .arg(Arg::from_usage("-d --directory=[DIRECTORY] 'event log files directory'")) .arg(Arg::from_usage("-s --statistics 'event statistics'")) .arg(Arg::from_usage("--credits 'Zachary Mathis, Akira Nishikawa'")) + .get_matches() +} + +fn is_test_mode() -> bool { + for i in std::env::args() { + if i == "--test" { + return true; + } + } + + return false; } #[derive(Debug, Clone)] @@ -70,10 +81,10 @@ impl EventKeyAliasConfig { } } -fn load_eventkey_alias() -> EventKeyAliasConfig { +fn load_eventkey_alias(path: &str) -> EventKeyAliasConfig { let mut config = EventKeyAliasConfig::new(); - let read_result = utils::read_csv("config/eventkey_alias.txt"); + let read_result = utils::read_csv(path); // eventkey_aliasが読み込めなかったらエラーで終了とする。 read_result.unwrap().into_iter().for_each(|line| { if line.len() != 2 { @@ -94,3 +105,27 @@ fn load_eventkey_alias() -> EventKeyAliasConfig { return config; } + +#[cfg(test)] +mod tests { + + use crate::detections::configs; + + #[test] + #[ignore] + fn singleton_read_and_write() { + let message = + "EventKeyAliasConfig { key_to_eventkey: {\"EventID\": \"Event.System.EventID\"} }"; + configs::CONFIG.write().unwrap().event_key_alias_config = + configs::load_eventkey_alias("test_files/config/eventkey_alias.txt"); + + let display = format!( + "{}", + format_args!( + "{:?}", + configs::CONFIG.write().unwrap().event_key_alias_config + ) + ); + assert_eq!(message, display); + } +} diff --git a/src/detections/print.rs b/src/detections/print.rs index 1adfd72f..37dd361f 100644 --- a/src/detections/print.rs +++ b/src/detections/print.rs @@ -70,7 +70,9 @@ impl Message { .take(target_length) .collect::(); - if let Some(array_str) = configs::singleton() + if let Some(array_str) = configs::CONFIG + .read() + .unwrap() .event_key_alias_config .get_event_key(target_str.to_string()) { diff --git a/src/detections/utils.rs b/src/detections/utils.rs index 704716a6..47b6ae85 100644 --- a/src/detections/utils.rs +++ b/src/detections/utils.rs @@ -91,9 +91,11 @@ pub fn get_event_value<'a>(key: &String, event_value: &'a Value) -> Option<&'a V if key.len() == 0 { return Option::None; } - - let alias_config = configs::singleton().event_key_alias_config; - let event_key = match alias_config.get_event_key(key.to_string()) { + let singleton = configs::CONFIG.read().unwrap(); + let event_key = match singleton + .event_key_alias_config + .get_event_key(key.to_string()) + { Some(alias_event_key) => alias_event_key, None => key, }; diff --git a/src/main.rs b/src/main.rs index a5812535..d467018f 100644 --- a/src/main.rs +++ b/src/main.rs @@ -8,12 +8,12 @@ use yamato_event_analyzer::detections::detection; use yamato_event_analyzer::omikuji::Omikuji; fn main() { - if let Some(filepath) = configs::singleton().args.value_of("filepath") { + if let Some(filepath) = configs::CONFIG.read().unwrap().args.value_of("filepath") { detect_files(vec![PathBuf::from(filepath)]); - } else if let Some(directory) = configs::singleton().args.value_of("directory") { + } else if let Some(directory) = configs::CONFIG.read().unwrap().args.value_of("directory") { let evtx_files = collect_evtxfiles(&directory); detect_files(evtx_files); - } else if configs::singleton().args.is_present("credits") { + } else if configs::CONFIG.read().unwrap().args.is_present("credits") { print_credits(); } } diff --git a/test_files/config/eventkey_alias.txt b/test_files/config/eventkey_alias.txt new file mode 100644 index 00000000..5080ad75 --- /dev/null +++ b/test_files/config/eventkey_alias.txt @@ -0,0 +1,2 @@ +alias,event_key +EventID,Event.System.EventID