Feature/changeSingleton#53,#65 (#67)

* change from singleton to global variable

* issue #65

* #65 test mode implemented

* cargo fmt --all
This commit is contained in:
nishikawaakira
2021-03-19 04:46:52 +09:00
committed by GitHub
parent f594bb1432
commit 7b0357b120
6 changed files with 86 additions and 41 deletions

View File

@@ -16,8 +16,12 @@ pub struct CsvFormat<'a> {
} }
pub fn after_fact() { pub fn after_fact() {
let mut target: Box<dyn io::Write> = let mut target: Box<dyn io::Write> = if let Some(csv_path) = configs::CONFIG
if let Some(csv_path) = configs::singleton().args.value_of("csv-timeline") { .read()
.unwrap()
.args
.value_of("csv-timeline")
{
match File::create(csv_path) { match File::create(csv_path) {
Ok(file) => Box::new(file), Ok(file) => Box::new(file),
Err(err) => { Err(err) => {
@@ -53,7 +57,7 @@ fn emit_csv(writer: &mut Box<dyn io::Write>) -> Result<(), Box<dyn Error>> {
} }
fn format_time(time: &DateTime<Utc>) -> String { fn format_time(time: &DateTime<Utc>) -> String {
if configs::singleton().args.is_present("utc") { if configs::CONFIG.read().unwrap().args.is_present("utc") {
format_rfc(time) format_rfc(time)
} else { } else {
format_rfc(&time.with_timezone(&Local)) format_rfc(&time.with_timezone(&Local))
@@ -64,7 +68,7 @@ fn format_rfc<Tz: TimeZone>(time: &DateTime<Tz>) -> String
where where
Tz::Offset: std::fmt::Display, 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(); return time.to_rfc2822();
} else { } else {
return time.to_rfc3339(); return time.to_rfc3339();

View File

@@ -1,32 +1,28 @@
use crate::detections::utils; use crate::detections::utils;
use clap::{App, AppSettings, Arg, ArgMatches}; use clap::{App, AppSettings, Arg, ArgMatches};
use lazy_static::lazy_static;
use std::collections::HashMap; use std::collections::HashMap;
use std::sync::Once; use std::sync::RwLock;
lazy_static! {
pub static ref CONFIG: RwLock<ConfigReader> = RwLock::new(ConfigReader::new());
}
#[derive(Clone)] #[derive(Clone)]
pub struct SingletonReader { pub struct ConfigReader {
pub args: ArgMatches<'static>, pub args: ArgMatches<'static>,
pub event_key_alias_config: EventKeyAliasConfig, pub event_key_alias_config: EventKeyAliasConfig,
} }
pub fn singleton() -> Box<SingletonReader> { impl ConfigReader {
static mut SINGLETON: Option<Box<SingletonReader>> = Option::None; pub fn new() -> Self {
static ONCE: Once = Once::new(); ConfigReader {
args: build_app(),
unsafe { event_key_alias_config: load_eventkey_alias("config/eventkey_alias.txt"),
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();
} }
} }
fn build_app() -> clap::App<'static, 'static> { fn build_app<'a>() -> ArgMatches<'a> {
let program = std::env::args() let program = std::env::args()
.nth(0) .nth(0)
.and_then(|s| { .and_then(|s| {
@@ -36,7 +32,11 @@ fn build_app() -> clap::App<'static, 'static> {
}) })
.unwrap(); .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!") .about("Yea! (Yamato Event Analyzer). Aiming to be the world's greatest Windows event log analysis tool!")
.version("0.0.1") .version("0.0.1")
.author("Author name <author@example.com>") .author("Author name <author@example.com>")
@@ -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("-d --directory=[DIRECTORY] 'event log files directory'"))
.arg(Arg::from_usage("-s --statistics 'event statistics'")) .arg(Arg::from_usage("-s --statistics 'event statistics'"))
.arg(Arg::from_usage("--credits 'Zachary Mathis, Akira Nishikawa'")) .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)] #[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 mut config = EventKeyAliasConfig::new();
let read_result = utils::read_csv("config/eventkey_alias.txt"); let read_result = utils::read_csv(path);
// eventkey_aliasが読み込めなかったらエラーで終了とする。 // eventkey_aliasが読み込めなかったらエラーで終了とする。
read_result.unwrap().into_iter().for_each(|line| { read_result.unwrap().into_iter().for_each(|line| {
if line.len() != 2 { if line.len() != 2 {
@@ -94,3 +105,27 @@ fn load_eventkey_alias() -> EventKeyAliasConfig {
return config; 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);
}
}

View File

@@ -70,7 +70,9 @@ impl Message {
.take(target_length) .take(target_length)
.collect::<String>(); .collect::<String>();
if let Some(array_str) = configs::singleton() if let Some(array_str) = configs::CONFIG
.read()
.unwrap()
.event_key_alias_config .event_key_alias_config
.get_event_key(target_str.to_string()) .get_event_key(target_str.to_string())
{ {

View File

@@ -91,9 +91,11 @@ pub fn get_event_value<'a>(key: &String, event_value: &'a Value) -> Option<&'a V
if key.len() == 0 { if key.len() == 0 {
return Option::None; return Option::None;
} }
let singleton = configs::CONFIG.read().unwrap();
let alias_config = configs::singleton().event_key_alias_config; let event_key = match singleton
let event_key = match alias_config.get_event_key(key.to_string()) { .event_key_alias_config
.get_event_key(key.to_string())
{
Some(alias_event_key) => alias_event_key, Some(alias_event_key) => alias_event_key,
None => key, None => key,
}; };

View File

@@ -8,12 +8,12 @@ use yamato_event_analyzer::detections::detection;
use yamato_event_analyzer::omikuji::Omikuji; use yamato_event_analyzer::omikuji::Omikuji;
fn main() { 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)]); 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); let evtx_files = collect_evtxfiles(&directory);
detect_files(evtx_files); detect_files(evtx_files);
} else if configs::singleton().args.is_present("credits") { } else if configs::CONFIG.read().unwrap().args.is_present("credits") {
print_credits(); print_credits();
} }
} }

View File

@@ -0,0 +1,2 @@
alias,event_key
EventID,Event.System.EventID