Feature/risk level condition#45 (#186)

* add risk level filter arguments #45

* fix default level in help #45

* add test yaml files #45

* refactoring and fix level argument usage.

* cargo fmt --all

Co-authored-by: ichiichi11 <takai.wa.hajime@gmail.com>

src/detections/configs.rs

@@ -5,6 +5,15 @@ use std::collections::HashMap;
 use std::sync::RwLock;
 lazy_static! {
     pub static ref CONFIG: RwLock<ConfigReader> = RwLock::new(ConfigReader::new());
+    pub static ref LEVELMAP: HashMap<String, u8> = {
+        let mut levelmap = HashMap::new();
+        levelmap.insert("INFO".to_owned(), 1);
+        levelmap.insert("LOW".to_owned(), 2);
+        levelmap.insert("MEDIUM".to_owned(), 3);
+        levelmap.insert("HIGH".to_owned(), 4);
+        levelmap.insert("CRITICAL".to_owned(), 5);
+        return levelmap;
+    };
 }
 
 #[derive(Clone)]
@@ -42,6 +51,7 @@ fn build_app<'a>() -> ArgMatches<'a> {
     --csv-timeline=[CSV_TIMELINE] 'Csv output timeline'
     --rfc-2822 'Output date and time in RFC 2822 format. Example: Mon, 07 Aug 2006 12:34:56 -0600'
     -l --lang=[LANG] 'Output language'
+    -L --level=[LEVEL] 'Specified execute rule level(default: INFO)'
     -u --utc 'Output time in UTC format(default: local time)'
     -d --directory=[DIRECTORY] 'Event log files directory'
     -s --statistics 'Prints statistics for event logs'

src/detections/detection.rs

@@ -46,10 +46,10 @@ impl Detection {
     }
 
     // ルールファイルをパースします。
-    pub fn parse_rule_files() -> Vec<RuleNode> {
+    pub fn parse_rule_files(level: String) -> Vec<RuleNode> {
         // ルールファイルのパースを実行
         let mut rulefile_loader = ParseYaml::new();
-        let resutl_readdir = rulefile_loader.read_dir(DIRPATH_RULES);
+        let resutl_readdir = rulefile_loader.read_dir(DIRPATH_RULES, &level);
         if resutl_readdir.is_err() {
             let stdout = std::io::stdout();
             let mut stdout = stdout.lock();

src/main.rs

@@ -104,11 +104,17 @@ fn print_credits() {
 }
 
 fn analysis_files(evtx_files: Vec<PathBuf>) {
-    let mut detection = detection::Detection::new(detection::Detection::parse_rule_files());
+    let level = configs::CONFIG
-
+        .read()
+        .unwrap()
+        .args
+        .value_of("level")
+        .unwrap_or("INFO")
+        .to_uppercase();
+    let rule_files = detection::Detection::parse_rule_files(level);
+    let mut detection = detection::Detection::new(rule_files);
     for evtx_file in evtx_files {
-        let ret = analysis_file(evtx_file, detection);
+        detection = analysis_file(evtx_file, detection);
-        detection = ret;
     }
 
     after_fact();

src/yaml.rs

@@ -1,12 +1,14 @@
 extern crate serde_derive;
 extern crate yaml_rust;
 
+use crate::detections::configs;
 use crate::detections::print::AlertMessage;
 use std::ffi::OsStr;
 use std::fs;
 use std::io;
 use std::io::{BufReader, Read};
 use std::path::{Path, PathBuf};
+use yaml_rust::Yaml;
 use yaml_rust::YamlLoader;
 
 pub struct ParseYaml {
@@ -31,53 +33,90 @@ impl ParseYaml {
         Ok(file_content)
     }
 
-    pub fn read_dir<P: AsRef<Path>>(&mut self, path: P) -> io::Result<String> {
+    pub fn read_dir<P: AsRef<Path>>(&mut self, path: P, level: &str) -> io::Result<String> {
-        Ok(fs::read_dir(path)?
-            .filter_map(|entry| {
-                let entry = entry.ok()?;
-                if entry.file_type().ok()?.is_file()
-                    && entry.path().extension().unwrap_or(OsStr::new("")) == "yml"
-                {
         let stdout = std::io::stdout();
         let mut stdout = stdout.lock();
-                    match self.read_file(entry.path()) {
+        let mut entries = fs::read_dir(path)?;
-                        Ok(s) => {
+        let yaml_docs = entries.try_fold(vec![], |mut ret, entry| {
-                            match YamlLoader::load_from_str(&s) {
+            let entry = entry?;
-                                Ok(docs) => {
+            // フォルダは再帰的に呼び出す。
-                                    for i in docs {
+            if entry.file_type()?.is_dir() {
-                                        // If there is no "enabled" it does not load
+                self.read_dir(entry.path(), level)?;
-                                        if i["ignore"].as_bool().unwrap_or(false) {
+                return io::Result::Ok(ret);
-                                            continue;
             }
-                                        &self
+            // ファイル以外は無視
-                                            .files
+            if !entry.file_type()?.is_file() {
-                                            .push((format!("{}", entry.path().display()), i));
+                return io::Result::Ok(ret);
             }
+
+            // 拡張子がymlでないファイルは無視
+            let path = entry.path();
+            if path.extension().unwrap_or(OsStr::new("")) != "yml" {
+                return io::Result::Ok(ret);
             }
-                                Err(e) => {
+
+            // 個別のファイルの読み込みは即終了としない。
+            let read_content = self.read_file(path);
+            if read_content.is_err() {
                 AlertMessage::alert(
                     &mut stdout,
-                                        format!("fail to read file\n{}\n{} ", s, e),
+                    format!(
-                                    )
+                        "fail to read file: {}\n{} ",
-                                    .ok();
+                        entry.path().display(),
+                        read_content.unwrap_err()
+                    ),
+                )?;
+                return io::Result::Ok(ret);
             }
-                            }
+
-                        }
+            // ここも個別のファイルの読み込みは即終了としない。
-                        Err(e) => {
+            let yaml_contents = YamlLoader::load_from_str(&read_content.unwrap());
+            if yaml_contents.is_err() {
                 AlertMessage::alert(
                     &mut stdout,
-                                format!("fail to read file: {}\n{} ", entry.path().display(), e),
+                    format!(
-                            )
+                        "fail to parse as yaml: {}\n{} ",
-                            .ok();
+                        entry.path().display(),
+                        yaml_contents.unwrap_err()
+                    ),
+                )?;
+                return io::Result::Ok(ret);
             }
-                    };
+
+            let yaml_contents = yaml_contents.unwrap().into_iter().map(|yaml_content| {
+                let filepath = format!("{}", entry.path().display());
+                return (filepath, yaml_content);
+            });
+            ret.extend(yaml_contents);
+            return io::Result::Ok(ret);
+        })?;
+
+        let files: Vec<(String, Yaml)> = yaml_docs
+            .into_iter()
+            .filter_map(|(filepath, yaml_doc)| {
+                // ignoreフラグがONになっているルールは無視する。
+                if yaml_doc["ignore"].as_bool().unwrap_or(false) {
+                    return Option::None;
                 }
-                if entry.file_type().ok()?.is_dir() {
+
-                    let _ = self.read_dir(entry.path());
+                // 指定されたレベルより低いルールは無視する
+                let doc_level = &yaml_doc["level"]
+                    .as_str()
+                    .unwrap_or("INFO")
+                    .to_string()
+                    .to_uppercase();
+                let doc_level_num = configs::LEVELMAP.get(doc_level).unwrap_or(&1);
+                let args_level_num = configs::LEVELMAP.get(level).unwrap_or(&1);
+                if doc_level_num < args_level_num {
+                    return Option::None;
                 }
-                Some("")
+
+                return Option::Some((filepath, yaml_doc));
             })
-            .collect())
+            .collect();
+        self.files.extend(files);
+
+        return io::Result::Ok(String::default());
     }
 }
 
@@ -91,7 +130,7 @@ mod tests {
     #[test]
     fn test_read_dir_yaml() {
         let mut yaml = yaml::ParseYaml::new();
-        &yaml.read_dir("test_files/rules/yaml/".to_string());
+        &yaml.read_dir("test_files/rules/yaml/".to_string(), &"".to_owned());
         assert_ne!(yaml.files.len(), 0);
     }
 
@@ -120,4 +159,50 @@ mod tests {
         let rule = YamlLoader::load_from_str(&ret);
         assert_eq!(rule.is_err(), true);
     }
+
+    #[test]
+    /// no specifed "level" arguments value is adapted default level(INFO)
+    fn test_default_level_read_yaml() {
+        let mut yaml = yaml::ParseYaml::new();
+        let path = Path::new("test_files/rules/level_yaml");
+        yaml.read_dir(path.to_path_buf(), &"").unwrap();
+        assert_eq!(yaml.files.len(), 5);
+    }
+
+    #[test]
+    fn test_info_level_read_yaml() {
+        let mut yaml = yaml::ParseYaml::new();
+        let path = Path::new("test_files/rules/level_yaml");
+        yaml.read_dir(path.to_path_buf(), &"INFO").unwrap();
+        assert_eq!(yaml.files.len(), 5);
+    }
+    #[test]
+    fn test_low_level_read_yaml() {
+        let mut yaml = yaml::ParseYaml::new();
+        let path = Path::new("test_files/rules/level_yaml");
+
+        yaml.read_dir(path.to_path_buf(), &"LOW").unwrap();
+        assert_eq!(yaml.files.len(), 4);
+    }
+    #[test]
+    fn test_medium_level_read_yaml() {
+        let mut yaml = yaml::ParseYaml::new();
+        let path = Path::new("test_files/rules/level_yaml");
+        yaml.read_dir(path.to_path_buf(), &"MEDIUM").unwrap();
+        assert_eq!(yaml.files.len(), 3);
+    }
+    #[test]
+    fn test_high_level_read_yaml() {
+        let mut yaml = yaml::ParseYaml::new();
+        let path = Path::new("test_files/rules/level_yaml");
+        yaml.read_dir(path.to_path_buf(), &"HIGH").unwrap();
+        assert_eq!(yaml.files.len(), 2);
+    }
+    #[test]
+    fn test_critical_level_read_yaml() {
+        let mut yaml = yaml::ParseYaml::new();
+        let path = Path::new("test_files/rules/level_yaml");
+        yaml.read_dir(path.to_path_buf(), &"CRITICAL").unwrap();
+        assert_eq!(yaml.files.len(), 1);
+    }
 }

test_files/rules/level_yaml/critical.yml

@@ -0,0 +1,19 @@
+title: Sysmon Check command lines
+description: hogehoge
+enabled: true
+author: Yea
+logsource: 
+    product: windows
+detection:
+    selection:
+        EventLog: Sysmon
+        EventID: 1
+        CommandLine: '*'
+    condition: selection
+falsepositives:
+    - unknown
+level: critical
+output: 'CommandLine=%CommandLine%¥nParentImage=%ParentImage%'
+creation_date: 2020/11/8
+updated_date: 2020/11/8
+

test_files/rules/level_yaml/high.yml

@@ -0,0 +1,19 @@
+title: Sysmon Check command lines
+description: hogehoge
+enabled: true
+author: Yea
+logsource: 
+    product: windows
+detection:
+    selection:
+        EventLog: Sysmon
+        EventID: 1
+        CommandLine: '*'
+    condition: selection
+falsepositives:
+    - unknown
+level: high
+output: 'CommandLine=%CommandLine%¥nParentImage=%ParentImage%'
+creation_date: 2020/11/8
+updated_date: 2020/11/8
+

test_files/rules/level_yaml/info.yml

@@ -0,0 +1,19 @@
+title: Sysmon Check command lines
+description: hogehoge
+enabled: true
+author: Yea
+logsource: 
+    product: windows
+detection:
+    selection:
+        EventLog: Sysmon
+        EventID: 1
+        CommandLine: '*'
+    condition: selection
+falsepositives:
+    - unknown
+level: info
+output: 'CommandLine=%CommandLine%¥nParentImage=%ParentImage%'
+creation_date: 2020/11/8
+updated_date: 2020/11/8
+

test_files/rules/level_yaml/low.yml

@@ -0,0 +1,19 @@
+title: Sysmon Check command lines
+description: hogehoge
+enabled: true
+author: Yea
+logsource: 
+    product: windows
+detection:
+    selection:
+        EventLog: Sysmon
+        EventID: 1
+        CommandLine: '*'
+    condition: selection
+falsepositives:
+    - unknown
+level: low
+output: 'CommandLine=%CommandLine%¥nParentImage=%ParentImage%'
+creation_date: 2020/11/8
+updated_date: 2020/11/8
+

test_files/rules/level_yaml/medium.yml

@@ -0,0 +1,19 @@
+title: Sysmon Check command lines
+description: hogehoge
+enabled: true
+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
+