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() {
let mut target: Box<dyn io::Write> =
if let Some(csv_path) = configs::singleton().args.value_of("csv-timeline") {
let mut target: Box<dyn io::Write> = 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) => {
@@ -53,7 +57,7 @@ fn emit_csv(writer: &mut Box<dyn io::Write>) -> Result<(), Box<dyn Error>> {
}
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)
} else {
format_rfc(&time.with_timezone(&Local))
@@ -64,7 +68,7 @@ fn format_rfc<Tz: TimeZone>(time: &DateTime<Tz>) -> 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();

View File

@@ -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<ConfigReader> = 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<SingletonReader> {
static mut SINGLETON: Option<Box<SingletonReader>> = 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 <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("-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);
}
}

View File

@@ -70,7 +70,9 @@ impl Message {
.take(target_length)
.collect::<String>();
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())
{

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 {
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,
};

View File

@@ -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();
}
}

View File

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