refactoring

2020-11-30 21:17:30 +09:00
parent cd8948de4c
commit 0e3d2ebaf4
4 changed files with 154 additions and 94 deletions
--- a/src/afterfact.rs
+++ b/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;
-use std::fs::{read_to_string, remove_file};
-use std::io::Read;
+#[cfg(test)]
+mod tests {

-#[test]
-fn test_emit_csv() {
-    {
-        let mut messages = print::MESSAGES.lock().unwrap();
-        let poke = Utc.ymd(1996, 2, 27).and_hms(1, 5, 1);
-        let json_str = r#"
+    use crate::afterfact::emit_csv;
+    use crate::detections::print;
+    use serde_json::Value;
+    use std::fs::{read_to_string, remove_file};
+
+    #[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") {
-        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());
+    }
 }
--- a/src/detections/detection.rs
+++ b/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();
-    }
 }
--- a/src/detections/print.rs
+++ b/src/detections/print.rs
@@ -25,39 +25,28 @@ impl Message {
    }

    /// メッセージを設定
-    pub fn insert(
-        &mut self,
-        mut time: Option<DateTime<Utc>>,
-        event_record: &Value,
-        output: Option<String>,
-    ) {
-        if Option::None == output {
+    pub fn insert(&mut self, event_record: &Value, output: String) {
+        if output.is_empty() {
            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 {
-            time = Option::Some(Utc.ymd(1970, 1, 1).and_hms(0, 0, 0));
-        }
-
-        match self.map.get_mut(&time.unwrap()) {
+        match self.map.get_mut(&time) {
            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 {
-        if Option::None == output {
-            return "".to_string();
-        }
-
-        let mut return_message: String = output.unwrap();
+    fn parse_message(&mut self, event_record: &Value, output: String) -> String {
+        let mut return_message: String = output;
        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]
-fn test_create_and_append_message() {
-    let mut message = Message::new();
-    let poke = Utc.ymd(1996, 2, 27).and_hms(1, 5, 1);
-    let taka = Utc.ymd(2000, 1, 21).and_hms(9, 6, 1);
+#[cfg(test)]
+mod tests {
+    use crate::detections::print::Message;
+    use serde_json::Value;

-    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(
-        Some(poke),
-        &event_record,
-        Some("CommandLine1: %CommandLine%".to_string()),
-    );
-    message.insert(
-        Some(poke),
-        &event_record,
-        Some("CommandLine2: %CommandLine%".to_string()),
-    );
-    message.insert(
-        Some(taka),
-        &event_record,
-        Some("CommandLine3: %CommandLine%".to_string()),
-    );
-    message.insert(
-        Option::None,
-        &event_record,
-        Some("CommandLine4: %CommandLine%".to_string()),
-    );
+        let json_str_2 = r##"
+    {
+        "Event": {
+            "EventData": {
+                "CommandLine": "hoge"
+            },
+            "System": {
+                "TimeCreated": {
+                    "#attributes":{
+                        "SystemTime": "1996-02-27T01:05:01Z"
+                    }
+                }
+            }
+        }
+    }
+    "##;
+        let event_record_2: Value = serde_json::from_str(json_str_2).unwrap();
+        message.insert(&event_record_2, "CommandLine2: %CommandLine%".to_string());

-    let display = format!("{}", format_args!("{:?}", message));
-    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);
+        let json_str_3 = r##"
+    {
+        "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);
+    }
 }
--- a/src/detections/rule.rs
+++ b/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() {
-            return Result::Ok(());
+        let mut errmsgs: Vec<String> = vec![];
+
+        // 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 {