From f58d5f316b66e52c6ad5e3627ef4cdd918d466d1 Mon Sep 17 00:00:00 2001
From: akiranishikawa <info@i-wind.jp>
Date: Mon, 7 Dec 2020 12:18:48 +0900
Subject: [PATCH] resolved #40

---
 rules/deep_blue_cli/powershell/4104.yml |  8 ++--
 src/yaml.rs                             | 56 +++++++++++++++++--------
 test_files/rules/yaml/1.yml             |  2 +-
 test_files/rules/yaml/error.yml         | 19 +++++++++
 4 files changed, 63 insertions(+), 22 deletions(-)
 create mode 100644 test_files/rules/yaml/error.yml

diff --git a/rules/deep_blue_cli/powershell/4104.yml b/rules/deep_blue_cli/powershell/4104.yml
index d9d3a8b4..78952873 100644
--- a/rules/deep_blue_cli/powershell/4104.yml
+++ b/rules/deep_blue_cli/powershell/4104.yml
@@ -6,14 +6,14 @@ logsource:
     product: windows
 detection:
     selection:
-        Channel: PowerShell
+        Channel: Microsoft-Windows-PowerShell/Operational
         EventID: 4104
         Path: null
-        ScriptBlockText: null
+        ScriptBlockText: '.+'
     # condition: selection
 falsepositives:
     - unknown
 level: medium
-output: 'command=%CommandLine%'
+output: 'command=%ScriptBlockText%'
 creation_date: 2020/11/8
-uodated_date: 2020/11/8
+updated_date: 2020/11/8
diff --git a/src/yaml.rs b/src/yaml.rs
index 781a5100..a3b7323d 100644
--- a/src/yaml.rs
+++ b/src/yaml.rs
@@ -36,18 +36,25 @@ impl ParseYaml {
                 if entry.file_type().ok()?.is_file() {
                     match self.read_file(entry.path()) {
                         Ok(s) => {
-                            let docs = YamlLoader::load_from_str(&s).unwrap();
-                            for i in docs {
-                                if i["enabled"].as_bool().unwrap() {
-                                    &self.files.push(i);
+                            match YamlLoader::load_from_str(&s) {
+                                Ok(docs) => {
+                                    for i in docs {
+                                        // If there is no "enabled" it does not load
+                                        if i["enabled"].as_bool().unwrap_or(false) {
+                                            &self.files.push(i);
+                                        }
+                                    }
                                 }
+                                Err(e) => eprintln!("fail to read file\n{}\n{} ", s, e),
                             }
                         }
-                        Err(e) => panic!("fail to read file: {}", e),
+                        Err(e) => {
+                            eprintln!("fail to read file: {}\n{} ", entry.path().display(), e)
+                        }
                     };
                 }
                 if entry.file_type().ok()?.is_dir() {
-                    self.read_dir(entry.path());
+                    let _ = self.read_dir(entry.path());
                 }
                 Some("")
             })
@@ -59,24 +66,39 @@ impl ParseYaml {
 mod tests {
 
     use crate::yaml;
+    use std::path::Path;
+    use yaml_rust::YamlLoader;
+
+    #[test]
+    fn test_read_dir_yaml() {
+        let mut yaml = yaml::ParseYaml::new();
+        &yaml.read_dir("test_files/rules/yaml/".to_string());
+        assert_ne!(yaml.files.len(), 0);
+    }
 
     #[test]
     fn test_read_yaml() {
-        let mut yaml = yaml::ParseYaml::new();
-        &yaml.read_dir("test_files/rules/yaml/".to_string());
-        for rule in yaml.files {
-            if rule["title"].as_str().unwrap() == "Sysmon Check command lines" {
+        let yaml = yaml::ParseYaml::new();
+        let path = Path::new("test_files/rules/yaml/1.yml");
+        let ret = yaml.read_file(path.to_path_buf()).unwrap();
+        let rule = YamlLoader::load_from_str(&ret).unwrap();
+        for i in rule {
+            if i["title"].as_str().unwrap() == "Sysmon Check command lines" {
                 assert_eq!(
                     "*",
-                    rule["detection"]["selection"]["CommandLine"]
-                        .as_str()
-                        .unwrap()
-                );
-                assert_eq!(
-                    1,
-                    rule["detection"]["selection"]["EventID"].as_i64().unwrap()
+                    i["detection"]["selection"]["CommandLine"].as_str().unwrap()
                 );
+                assert_eq!(1, i["detection"]["selection"]["EventID"].as_i64().unwrap());
             }
         }
     }
+
+    #[test]
+    fn test_failed_read_yaml() {
+        let yaml = yaml::ParseYaml::new();
+        let path = Path::new("test_files/rules/yaml/error.yml");
+        let ret = yaml.read_file(path.to_path_buf()).unwrap();
+        let rule = YamlLoader::load_from_str(&ret);
+        assert_eq!(rule.is_err(), true);
+    }
 }
diff --git a/test_files/rules/yaml/1.yml b/test_files/rules/yaml/1.yml
index 26e06dbb..5f844d26 100644
--- a/test_files/rules/yaml/1.yml
+++ b/test_files/rules/yaml/1.yml
@@ -15,5 +15,5 @@ falsepositives:
 level: medium
 output: 'CommandLine=%CommandLine%¥nParentImage=%ParentImage%'
 creation_date: 2020/11/8
-uodated_date: 2020/11/8
+updated_date: 2020/11/8
 
diff --git a/test_files/rules/yaml/error.yml b/test_files/rules/yaml/error.yml
new file mode 100644
index 00000000..b897264f
--- /dev/null
+++ b/test_files/rules/yaml/error.yml
@@ -0,0 +1,19 @@
+title: Sysmon Check command lines
+description: hogehoge
+enabled: truea
+author: Yea
+logsource: 
+    product: windows
+detection:
+    selection:
+        EventLog: Sysmon
+        EventID: 1
+        CommandLine: %
+    condition: selection
+falsepositives:
+    - unknown
+level: medium
+output: 'CommandLine=%CommandLine%¥nParentImage=%ParentImage%'
+creation_date: 2020/11/8
+updated_date: 2020/11/8
+