Update: release csv-timeline function

This commit is contained in:
itiB
2020-12-03 23:26:17 +09:00
parent a30707d736
commit 12110a4687
2 changed files with 49 additions and 44 deletions

View File

@@ -1,8 +1,10 @@
use crate::detections::configs; use crate::detections::configs;
use crate::detections::print; use crate::detections::print;
use chrono::{DateTime, Utc}; use chrono::{DateTime, TimeZone, Utc};
use serde::Serialize; use serde::Serialize;
use std::error::Error; use std::error::Error;
use std::fs::File;
use std::io;
use std::process; use std::process;
#[derive(Debug, Serialize)] #[derive(Debug, Serialize)]
@@ -13,16 +15,27 @@ pub struct CsvFormat<'a> {
} }
pub fn after_fact() { pub fn after_fact() {
let mut target: Box<dyn io::Write> =
if let Some(csv_path) = configs::singleton().args.value_of("csv-timeline") { if let Some(csv_path) = configs::singleton().args.value_of("csv-timeline") {
if let Err(err) = emit_csv(csv_path) { match File::create(csv_path) {
println!("{}", err); Ok(file) => Box::new(file),
Err(err) => {
println!("Failed to open file. {}", err);
process::exit(1); 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<dyn Error>> { fn emit_csv(writer: &mut Box<dyn io::Write>) -> Result<(), Box<dyn Error>> {
let mut wtr = csv::Writer::from_path(path)?; let mut wtr = csv::WriterBuilder::new().from_writer(writer);
let messages = print::MESSAGES.lock().unwrap(); let messages = print::MESSAGES.lock().unwrap();
for (time, texts) in messages.iter() { for (time, texts) in messages.iter() {
@@ -37,19 +50,14 @@ fn emit_csv(path: &str) -> Result<(), Box<dyn Error>> {
Ok(()) Ok(())
} }
#[cfg(test)] #[test]
mod tests { fn test_emit_csv() {
use crate::afterfact::emit_csv;
use crate::detections::print;
use serde_json::Value; use serde_json::Value;
use std::fs::{read_to_string, remove_file}; use std::fs::{read_to_string, remove_file};
#[test]
fn test_emit_csv() {
{ {
let mut messages = print::MESSAGES.lock().unwrap(); let mut messages = print::MESSAGES.lock().unwrap();
let json_str = r##"
let val = r##"
{ {
"Event": { "Event": {
"EventData": { "EventData": {
@@ -65,16 +73,17 @@ mod tests {
} }
} }
"##; "##;
let event_record: Value = serde_json::from_str(json_str).unwrap(); let event: Value = serde_json::from_str(val).unwrap();
messages.insert(&event, "pokepoke".to_string());
messages.insert(&event_record, "pokepoke".to_string());
} }
let expect = "Time,Message let expect = "Time,Message
1996-02-27T01:05:01Z,pokepoke 1996-02-27T01:05:01Z,pokepoke
"; ";
assert!(emit_csv(&"./test_emit_csv.csv".to_string()).is_ok()); let mut file: Box<dyn io::Write> =
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") { match read_to_string("./test_emit_csv.csv") {
Err(_) => panic!("Failed to open file"), Err(_) => panic!("Failed to open file"),
@@ -82,5 +91,4 @@ mod tests {
}; };
assert!(remove_file("./test_emit_csv.csv").is_ok()); assert!(remove_file("./test_emit_csv.csv").is_ok());
}
} }

View File

@@ -1,7 +1,7 @@
extern crate chrono; extern crate chrono;
extern crate csv; extern crate csv;
use crate::detections::print::Message; use crate::detections::print::MESSAGES;
use crate::detections::rule; use crate::detections::rule;
use crate::detections::rule::RuleNode; use crate::detections::rule::RuleNode;
use crate::yaml::ParseYaml; use crate::yaml::ParseYaml;
@@ -81,7 +81,7 @@ impl Detection {
.collect(); .collect();
// selection rule files and collect message // selection rule files and collect message
let mut message = Message::new(); let mut message = MESSAGES.lock().unwrap();
selection_rules.iter_mut().for_each(|rule| { selection_rules.iter_mut().for_each(|rule| {
event_records.iter().for_each(|event_record| { event_records.iter().for_each(|event_record| {
if !rule.select(event_record) { if !rule.select(event_record) {
@@ -94,8 +94,5 @@ impl Detection {
) )
}); });
}); });
// output message
message.print();
} }
} }