diff --git a/src/afterfact.rs b/src/afterfact.rs index 002d3d1a..362c7a81 100644 --- a/src/afterfact.rs +++ b/src/afterfact.rs @@ -1,8 +1,10 @@ use crate::detections::configs; use crate::detections::print; -use chrono::{DateTime, Utc}; +use chrono::{DateTime, TimeZone, Utc}; use serde::Serialize; use std::error::Error; +use std::fs::File; +use std::io; use std::process; #[derive(Debug, Serialize)] @@ -13,16 +15,27 @@ pub struct CsvFormat<'a> { } pub fn after_fact() { - if let Some(csv_path) = configs::singleton().args.value_of("csv-timeline") { - if let Err(err) = emit_csv(csv_path) { - println!("{}", err); - process::exit(1); - } + 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); + } + } + } else { + Box::new(io::stdout()) + }; + + if let Err(err) = emit_csv(&mut target) { + println!("Failed to write CSV. {}", err); + process::exit(1); } } -fn emit_csv(path: &str) -> Result<(), Box> { - let mut wtr = csv::Writer::from_path(path)?; +fn emit_csv(writer: &mut Box) -> Result<(), Box> { + let mut wtr = csv::WriterBuilder::new().from_writer(writer); let messages = print::MESSAGES.lock().unwrap(); for (time, texts) in messages.iter() { @@ -37,50 +50,45 @@ fn emit_csv(path: &str) -> Result<(), Box> { Ok(()) } -#[cfg(test)] -mod tests { - - use crate::afterfact::emit_csv; - use crate::detections::print; +#[test] +fn test_emit_csv() { use serde_json::Value; use std::fs::{read_to_string, remove_file}; + { + let mut messages = print::MESSAGES.lock().unwrap(); - #[test] - fn test_emit_csv() { + let val = r##" { - let mut messages = print::MESSAGES.lock().unwrap(); - let json_str = r##" - { - "Event": { - "EventData": { - "CommandLine": "hoge" - }, - "System": { - "TimeCreated": { - "#attributes":{ - "SystemTime": "1996-02-27T01:05:01Z" - } + "Event": { + "EventData": { + "CommandLine": "hoge" + }, + "System": { + "TimeCreated": { + "#attributes":{ + "SystemTime": "1996-02-27T01:05:01Z" } } } } - "##; - let event_record: Value = serde_json::from_str(json_str).unwrap(); - - messages.insert(&event_record, "pokepoke".to_string()); } + "##; + let event: Value = serde_json::from_str(val).unwrap(); + messages.insert(&event, "pokepoke".to_string()); + } - let expect = "Time,Message + let expect = "Time,Message 1996-02-27T01:05:01Z,pokepoke "; - assert!(emit_csv(&"./test_emit_csv.csv".to_string()).is_ok()); + let mut file: Box = + Box::new(File::create("./test_emit_csv.csv".to_string()).unwrap()); + assert!(emit_csv(&mut file).is_ok()); - match read_to_string("./test_emit_csv.csv") { - Err(_) => panic!("Failed to open file"), - Ok(s) => assert_eq!(s, expect), - }; + match read_to_string("./test_emit_csv.csv") { + Err(_) => panic!("Failed to open file"), + Ok(s) => assert_eq!(s, expect), + }; - assert!(remove_file("./test_emit_csv.csv").is_ok()); - } + assert!(remove_file("./test_emit_csv.csv").is_ok()); } diff --git a/src/detections/detection.rs b/src/detections/detection.rs index 20977de4..e5d215ad 100644 --- a/src/detections/detection.rs +++ b/src/detections/detection.rs @@ -1,7 +1,7 @@ extern crate chrono; extern crate csv; -use crate::detections::print::Message; +use crate::detections::print::MESSAGES; use crate::detections::rule; use crate::detections::rule::RuleNode; use crate::yaml::ParseYaml; @@ -81,7 +81,7 @@ impl Detection { .collect(); // selection rule files and collect message - let mut message = Message::new(); + let mut message = MESSAGES.lock().unwrap(); selection_rules.iter_mut().for_each(|rule| { event_records.iter().for_each(|event_record| { if !rule.select(event_record) { @@ -94,8 +94,5 @@ impl Detection { ) }); }); - - // output message - message.print(); } }