refactoring

src/afterfact.rs

@@ -1,6 +1,6 @@
 use crate::detections::configs;
 use crate::detections::print;
-use chrono::{DateTime, TimeZone, Utc};
+use chrono::{DateTime, Utc};
 use serde::Serialize;
 use std::error::Error;
 use std::process;
@@ -37,39 +37,50 @@ fn emit_csv(path: &str) -> Result<(), Box<dyn Error>> {
     Ok(())
 }
 
-use serde_json::Value;
+#[cfg(test)]
-use std::fs::{read_to_string, remove_file};
+mod tests {
-use std::io::Read;
 
-#[test]
+    use crate::afterfact::emit_csv;
-fn test_emit_csv() {
+    use crate::detections::print;
-    {
+    use serde_json::Value;
-        let mut messages = print::MESSAGES.lock().unwrap();
+    use std::fs::{read_to_string, remove_file};
-        let poke = Utc.ymd(1996, 2, 27).and_hms(1, 5, 1);
+
-        let json_str = r#"
+    #[test]
+    fn test_emit_csv() {
+        {
+            let mut messages = print::MESSAGES.lock().unwrap();
+            let json_str = r##"
             {
                 "Event": {
                     "EventData": {
                         "CommandLine": "hoge"
+                    },
+                    "System": {
+                        "TimeCreated": {
+                            "#attributes":{
+                                "SystemTime": "1996-02-27T01:05:01Z"
+                            }
+                        }
                     }
                 }
             }
-        "#;
+        "##;
-        let event_record: Value = serde_json::from_str(json_str).unwrap();
+            let event_record: Value = serde_json::from_str(json_str).unwrap();
 
-        messages.insert(Some(poke), &event_record, Some("pokepoke".to_string()));
+            messages.insert(&event_record, "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());
+        assert!(emit_csv(&"./test_emit_csv.csv".to_string()).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"),
-        Ok(s) => assert_eq!(s, expect),
+            Ok(s) => assert_eq!(s, expect),
-    };
+        };
 
-    assert!(remove_file("./test_emit_csv.csv").is_ok());
+        assert!(remove_file("./test_emit_csv.csv").is_ok());
+    }
 }

src/detections/detection.rs

@@ -88,17 +88,9 @@ impl Detection {
                     return;
                 }
 
-                let event_time = Detection::get_event_time(event_record);
-                let utc_event_time = event_time
-                    .and_then(|datetime| {
-                        let utc = Utc.from_local_datetime(&datetime.naive_utc()).unwrap();
-                        return Option::Some(utc);
-                    })
-                    .or(Option::None);
                 message.insert(
-                    utc_event_time,
                     event_record,
-                    Some(rule.yaml["output"].as_str().unwrap().to_string()),
+                    rule.yaml["output"].as_str().unwrap_or("").to_string(),
                 )
             });
         });
@@ -106,15 +98,4 @@ impl Detection {
         // output message
         message.print();
     }
-
-    fn get_event_time(event_record: &Value) -> Option<DateTime<FixedOffset>> {
-        let system_time =
-            &event_record["Event"]["System"]["TimeCreated"]["#attributes"]["SystemTime"];
-        let system_time_str = system_time.as_str().unwrap_or("");
-        if system_time_str.is_empty() {
-            return Option::None;
-        }
-
-        return DateTime::parse_from_rfc3339(system_time_str).ok();
-    }
 }

src/detections/print.rs

@@ -25,39 +25,28 @@ impl Message {
     }
 
     /// メッセージを設定
-    pub fn insert(
+    pub fn insert(&mut self, event_record: &Value, output: String) {
-        &mut self,
+        if output.is_empty() {
-        mut time: Option<DateTime<Utc>>,
-        event_record: &Value,
-        output: Option<String>,
-    ) {
-        if Option::None == output {
             return;
         }
 
         let message = &self.parse_message(event_record, output);
+        let default_time = Utc.ymd(1970, 1, 1).and_hms(0, 0, 0);
+        let time = Message::get_event_time(event_record).unwrap_or(default_time);
 
-        if Option::None == time {
+        match self.map.get_mut(&time) {
-            time = Option::Some(Utc.ymd(1970, 1, 1).and_hms(0, 0, 0));
-        }
-
-        match self.map.get_mut(&time.unwrap()) {
             Some(v) => {
                 v.push(message.to_string());
             }
             None => {
                 let m = vec![message.to_string(); 1];
-                self.map.insert(time.unwrap(), m);
+                self.map.insert(time, m);
             }
         }
     }
 
-    fn parse_message(&mut self, event_record: &Value, output: Option<String>) -> String {
+    fn parse_message(&mut self, event_record: &Value, output: String) -> String {
-        if Option::None == output {
+        let mut return_message: String = output;
-            return "".to_string();
-        }
-
-        let mut return_message: String = output.unwrap();
         let mut hash_map: HashMap<String, String> = HashMap::new();
         let re = Regex::new(r"%[a-zA-Z0-9-_]+%").unwrap();
         for caps in re.captures_iter(&return_message) {
@@ -119,47 +108,110 @@ impl Message {
     pub fn iter(&self) -> &BTreeMap<DateTime<Utc>, Vec<String>> {
         &self.map
     }
+
+    fn get_event_time(event_record: &Value) -> Option<DateTime<Utc>> {
+        let system_time =
+            &event_record["Event"]["System"]["TimeCreated"]["#attributes"]["SystemTime"];
+        let system_time_str = system_time.as_str().unwrap_or("");
+        if system_time_str.is_empty() {
+            return Option::None;
+        }
+
+        let rfc3339_time = DateTime::parse_from_rfc3339(system_time_str);
+        if rfc3339_time.is_err() {
+            return Option::None;
+        }
+        let datetime = Utc
+            .from_local_datetime(&rfc3339_time.unwrap().naive_utc())
+            .single();
+        if datetime.is_none() {
+            return Option::None;
+        } else {
+            return Option::Some(datetime.unwrap());
+        }
+    }
 }
 
-#[test]
+#[cfg(test)]
-fn test_create_and_append_message() {
+mod tests {
-    let mut message = Message::new();
+    use crate::detections::print::Message;
-    let poke = Utc.ymd(1996, 2, 27).and_hms(1, 5, 1);
+    use serde_json::Value;
-    let taka = Utc.ymd(2000, 1, 21).and_hms(9, 6, 1);
 
-    let json_str = r#"
+    #[test]
+    fn test_create_and_append_message() {
+        let mut message = Message::new();
+        let json_str_1 = r##"
         {
             "Event": {
                 "EventData": {
                     "CommandLine": "hoge"
+                },
+                "System": {
+                    "TimeCreated": {
+                        "#attributes":{
+                            "SystemTime": "1996-02-27T01:05:01Z"
+                        }
+                    }
                 }
             }
         }
-    "#;
+    "##;
-    let event_record: Value = serde_json::from_str(json_str).unwrap();
+        let event_record_1: Value = serde_json::from_str(json_str_1).unwrap();
+        message.insert(&event_record_1, "CommandLine1: %CommandLine%".to_string());
 
-    message.insert(
+        let json_str_2 = r##"
-        Some(poke),
+    {
-        &event_record,
+        "Event": {
-        Some("CommandLine1: %CommandLine%".to_string()),
+            "EventData": {
-    );
+                "CommandLine": "hoge"
-    message.insert(
+            },
-        Some(poke),
+            "System": {
-        &event_record,
+                "TimeCreated": {
-        Some("CommandLine2: %CommandLine%".to_string()),
+                    "#attributes":{
-    );
+                        "SystemTime": "1996-02-27T01:05:01Z"
-    message.insert(
+                    }
-        Some(taka),
+                }
-        &event_record,
+            }
-        Some("CommandLine3: %CommandLine%".to_string()),
+        }
-    );
+    }
-    message.insert(
+    "##;
-        Option::None,
+        let event_record_2: Value = serde_json::from_str(json_str_2).unwrap();
-        &event_record,
+        message.insert(&event_record_2, "CommandLine2: %CommandLine%".to_string());
-        Some("CommandLine4: %CommandLine%".to_string()),
-    );
 
-    let display = format!("{}", format_args!("{:?}", message));
+        let json_str_3 = r##"
-    let expect = "Message { map: {1970-01-01T00:00:00Z: [\"CommandLine4: hoge\"], 1996-02-27T01:05:01Z: [\"CommandLine1: hoge\", \"CommandLine2: hoge\"], 2000-01-21T09:06:01Z: [\"CommandLine3: hoge\"]} }";
+    {
-    assert_eq!(display, expect);
+        "Event": {
+            "EventData": {
+                "CommandLine": "hoge"
+            },
+            "System": {
+                "TimeCreated": {
+                    "#attributes":{
+                        "SystemTime": "2000-01-21T09:06:01Z"
+                    }
+                }
+            }
+        }
+    }
+    "##;
+        let event_record_3: Value = serde_json::from_str(json_str_3).unwrap();
+        message.insert(&event_record_3, "CommandLine3: %CommandLine%".to_string());
+
+        let json_str_4 = r##"
+    {
+        "Event": {
+            "EventData": {
+                "CommandLine": "hoge"
+            }
+        }
+    }
+    "##;
+        let event_record_4: Value = serde_json::from_str(json_str_4).unwrap();
+        message.insert(&event_record_4, "CommandLine4: %CommandLine%".to_string());
+
+        let display = format!("{}", format_args!("{:?}", message));
+        println!("display::::{}", display);
+        let expect = "Message { map: {1970-01-01T00:00:00Z: [\"CommandLine4: hoge\"], 1996-02-27T01:05:01Z: [\"CommandLine1: hoge\", \"CommandLine2: hoge\"], 2000-01-21T09:06:01Z: [\"CommandLine3: hoge\"]} }";
+        assert_eq!(display, expect);
+    }
 }

src/detections/rule.rs

@@ -82,11 +82,27 @@ pub struct RuleNode {
 
 impl RuleNode {
     pub fn init(&mut self) -> Result<(), Vec<String>> {
-        if self.detection.is_none() {
+        let mut errmsgs: Vec<String> = vec![];
-            return Result::Ok(());
+
+        // field check
+        if self.yaml["output"].as_str().is_none() {
+            errmsgs.push("Cannot find required key. key:output".to_string());
         }
 
-        return self.detection.as_mut().unwrap().init();
+        // detection node initialization
+        self.detection.as_mut().and_then(|detection| {
+            let detection_result = detection.init();
+            if detection_result.is_err() {
+                errmsgs.extend(detection_result.unwrap_err());
+            }
+            return Option::Some(detection);
+        });
+
+        if errmsgs.is_empty() {
+            return Result::Ok(());
+        } else {
+            return Result::Err(errmsgs);
+        }
     }
 
     pub fn select(&mut self, event_record: &Value) -> bool {