Merge pull request #43 from YamatoSecurity/feature/emit_csv2

Update: release csv-timeline function

src/afterfact.rs

@@ -1,34 +1,47 @@
 use crate::detections::configs;
 use crate::detections::print;
-use chrono::{DateTime, Utc};
+use chrono::{DateTime, Local, TimeZone, Utc};
 use serde::Serialize;
 use std::error::Error;
+use std::fs::File;
+use std::io;
 use std::process;
 
 #[derive(Debug, Serialize)]
 #[serde(rename_all = "PascalCase")]
 pub struct CsvFormat<'a> {
-    time: DateTime<Utc>,
+    time: &'a str,
     message: &'a str,
 }
 
 pub fn after_fact() {
-    if let Some(csv_path) = configs::singleton().args.value_of("csv-timeline") {
+    let mut target: Box<dyn io::Write> =
-        if let Err(err) = emit_csv(csv_path) {
+        if let Some(csv_path) = configs::singleton().args.value_of("csv-timeline") {
-            println!("{}", err);
+            match File::create(csv_path) {
-            process::exit(1);
+                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<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();
 
     for (time, texts) in messages.iter() {
         for text in texts {
             wtr.serialize(CsvFormat {
-                time: *time,
+                time: &format_time(time),
                 message: text,
             })?;
         }
@@ -37,50 +50,65 @@ fn emit_csv(path: &str) -> Result<(), Box<dyn Error>> {
     Ok(())
 }
 
-#[cfg(test)]
+fn format_time(time: &DateTime<Utc>) -> String {
-mod tests {
+    if configs::singleton().args.is_present("utc") {
+        format_rfc(time)
+    } else {
+        format_rfc(&time.with_timezone(&Local))
+    }
+}
 
-    use crate::afterfact::emit_csv;
+fn format_rfc<Tz: TimeZone>(time: &DateTime<Tz>) -> String
-    use crate::detections::print;
+where
+    Tz::Offset: std::fmt::Display,
+{
+    if configs::singleton().args.is_present("rfc-2822") {
+        return time.to_rfc2822();
+    } else {
+        return time.to_rfc3339();
+    }
+}
+
+#[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]
+        let val = r##"
-    fn test_emit_csv() {
         {
-            let mut messages = print::MESSAGES.lock().unwrap();
+            "Event": {
-            let json_str = r##"
+                "EventData": {
-            {
+                    "CommandLine": "hoge"
-                "Event": {
+                },
-                    "EventData": {
+                "System": {
-                        "CommandLine": "hoge"
+                    "TimeCreated": {
-                    },
+                        "#attributes":{
-                    "System": {
+                            "SystemTime": "1996-02-27T01:05:01Z"
-                        "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 expect = "Time,Message
+        let event: Value = serde_json::from_str(val).unwrap();
-1996-02-27T01:05:01Z,pokepoke
+        messages.insert(&event, "pokepoke".to_string());
-";
-
-        assert!(emit_csv(&"./test_emit_csv.csv".to_string()).is_ok());
-
-        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());
     }
+
+    let expect = "Time,Message
+1996-02-2";
+
+    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") {
+        Err(_) => panic!("Failed to open file"),
+        Ok(s) => {
+            assert_eq!(&s[0..22], expect);
+        }
+    };
+
+    assert!(remove_file("./test_emit_csv.csv").is_ok());
 }

src/detections/configs.rs

@@ -45,11 +45,11 @@ fn build_app() -> clap::App<'static, 'static> {
         .arg(Arg::from_usage("--attackhunt=[ATTACK_HUNT] 'Attack Hunt'"))
         .arg(Arg::from_usage("--csv-timeline=[CSV_TIMELINE] 'csv output timeline'"))
         .arg(Arg::from_usage("--human-readable-timeline=[HUMAN_READABLE_TIMELINE] 'human readable timeline'"))
+        .arg(Arg::from_usage("--rfc-2822 'output date and time in RFC 2822 format. Example: Mon, 07 Aug 2006 12:34:56 -0600'"))
         .arg(Arg::from_usage("-l --lang=[LANG] 'output language'"))
-        .arg(Arg::from_usage("-t --timezone=[TIMEZONE] 'timezone setting'"))
+        .arg(Arg::from_usage("-u --utc 'output time in UTC format(default: local time)'"))
         .arg(Arg::from_usage("-d --directory=[DIRECTORY] 'event log files directory'"))
         .arg(Arg::from_usage("-s --statistics 'event statistics'"))
-        .arg(Arg::from_usage("-u --update 'signature update'"))
         .arg(Arg::from_usage("--credits 'Zachary Mathis, Akira Nishikawa'"))
 }
 

src/detections/detection.rs

@@ -1,15 +1,13 @@
 extern crate csv;
 
-use std::path::PathBuf;
+use crate::detections::print::MESSAGES;
-
-use crate::detections::print::Message;
 use crate::detections::rule;
 use crate::detections::rule::RuleNode;
 use crate::yaml::ParseYaml;
-
 use evtx::err;
 use evtx::{EvtxParser, SerializedEvtxRecord};
 use serde_json::{Error, Value};
+use std::path::PathBuf;
 
 const DIRPATH_RULES: &str = "rules";
 
@@ -37,7 +35,7 @@ impl Detection {
         let evtx_records = self.serialize_evtx_to_jsons(evtx_files);
 
         // select rule files and collect message
-        let mut message = Message::new();
+        let mut message = MESSAGES.lock().unwrap();
         selection_rules.iter_mut().for_each(|rule| {
             evtx_records.iter().for_each(|event_record| {
                 if !rule.select(event_record) {

src/main.rs

@@ -13,13 +13,9 @@ fn main() {
     } else if let Some(directory) = configs::singleton().args.value_of("directory") {
         let evtx_files = collect_evtxfiles(&directory);
         detect_files(evtx_files);
-    }
+    } else if configs::singleton().args.is_present("credits") {
-
-    if configs::singleton().args.is_present("credits") {
         print_credits();
     }
-
-    after_fact();
 }
 
 fn collect_evtxfiles(dirpath: &str) -> Vec<PathBuf> {
@@ -63,6 +59,8 @@ fn print_credits() {
 fn detect_files(evtx_files: Vec<PathBuf>) {
     let mut detection = detection::Detection::new();
     &detection.start(evtx_files);
+
+    after_fact();
 }
 
 fn _output_with_omikuji(omikuji: Omikuji) {