ルール場所指定オプションでファイルを扱えるようにする (#364)

* add only rule file path in --rules

* add error handling for metadata

* refactor

* add test

* rename test function

src/detections/configs.rs

@@ -55,7 +55,7 @@ fn build_app<'a>() -> ArgMatches<'a> {
 
     let usages = "-d --directory=[DIRECTORY] 'Directory of multiple .evtx files'
     -f --filepath=[FILEPATH] 'File path to one .evtx file'
-    -r --rules=[RULEDIRECTORY] 'Rule file directory (default: ./rules)'
+    -r --rules=[RULEDIRECTORY/RULEFILE] 'Rule file or directory (default: ./rules)'
     -o --output=[CSV_TIMELINE] 'Save the timeline in CSV format. Example: results.csv'
     -v --verbose 'Output verbose information'
     -D --enable-deprecated-rules 'Enable sigma rules marked as deprecated'

src/yaml.rs

@@ -52,31 +52,41 @@ impl ParseYaml {
         level: &str,
         exclude_ids: &RuleExclude,
     ) -> io::Result<String> {
-        let mut entries = fs::read_dir(path)?;
-        let yaml_docs = entries.try_fold(vec![], |mut ret, entry| {
-            let entry = entry?;
-            // フォルダは再帰的に呼び出す。
-            if entry.file_type()?.is_dir() {
-                self.read_dir(entry.path(), level, exclude_ids)?;
-                return io::Result::Ok(ret);
+        let metadata = fs::metadata(path.as_ref());
+        if metadata.is_err() {
+            let errmsg = format!(
+                "fail to read metadata of file: {}",
+                path.as_ref().to_path_buf().display(),
+            );
+            if configs::CONFIG.read().unwrap().args.is_present("verbose") {
+                AlertMessage::alert(&mut BufWriter::new(std::io::stderr().lock()), &errmsg)?;
             }
-            // ファイル以外は無視
-            if !entry.file_type()?.is_file() {
-                return io::Result::Ok(ret);
+            if !*QUIET_ERRORS_FLAG {
+                ERROR_LOG_STACK
+                    .lock()
+                    .unwrap()
+                    .push(format!("[ERROR] {}", errmsg));
             }
-
+        }
+        let mut yaml_docs = vec![];
+        if metadata.unwrap().file_type().is_file() {
             // 拡張子がymlでないファイルは無視
-            let path = entry.path();
-            if path.extension().unwrap_or(OsStr::new("")) != "yml" {
-                return io::Result::Ok(ret);
+            if path
+                .as_ref()
+                .to_path_buf()
+                .extension()
+                .unwrap_or(OsStr::new(""))
+                != "yml"
+            {
+                return io::Result::Ok(String::default());
             }
 
             // 個別のファイルの読み込みは即終了としない。
-            let read_content = self.read_file(path);
+            let read_content = self.read_file(path.as_ref().to_path_buf());
             if read_content.is_err() {
                 let errmsg = format!(
                     "fail to read file: {}\n{} ",
-                    entry.path().display(),
+                    path.as_ref().to_path_buf().display(),
                     read_content.unwrap_err()
                 );
                 if configs::CONFIG.read().unwrap().args.is_present("verbose") {
@@ -89,7 +99,7 @@ impl ParseYaml {
                         .push(format!("[WARN] {}", errmsg));
                 }
                 self.errorrule_count += 1;
-                return io::Result::Ok(ret);
+                return io::Result::Ok(String::default());
             }
 
             // ここも個別のファイルの読み込みは即終了としない。
@@ -97,7 +107,7 @@ impl ParseYaml {
             if yaml_contents.is_err() {
                 let errmsg = format!(
                     "Failed to parse yml: {}\n{} ",
-                    entry.path().display(),
+                    path.as_ref().to_path_buf().display(),
                     yaml_contents.unwrap_err()
                 );
                 if configs::CONFIG.read().unwrap().args.is_present("verbose") {
@@ -110,16 +120,83 @@ impl ParseYaml {
                         .push(format!("[WARN] {}", errmsg));
                 }
                 self.errorrule_count += 1;
-                return io::Result::Ok(ret);
+                return io::Result::Ok(String::default());
             }
 
-            let yaml_contents = yaml_contents.unwrap().into_iter().map(|yaml_content| {
-                let filepath = format!("{}", entry.path().display());
+            yaml_docs.extend(yaml_contents.unwrap().into_iter().map(|yaml_content| {
+                let filepath = format!("{}", path.as_ref().to_path_buf().display());
                 return (filepath, yaml_content);
-            });
-            ret.extend(yaml_contents);
-            return io::Result::Ok(ret);
-        })?;
+            }));
+        } else {
+            let mut entries = fs::read_dir(path)?;
+            yaml_docs = entries.try_fold(vec![], |mut ret, entry| {
+                let entry = entry?;
+                // フォルダは再帰的に呼び出す。
+                if entry.file_type()?.is_dir() {
+                    self.read_dir(entry.path(), level, exclude_ids)?;
+                    return io::Result::Ok(ret);
+                }
+                // ファイル以外は無視
+                if !entry.file_type()?.is_file() {
+                    return io::Result::Ok(ret);
+                }
+
+                // 拡張子がymlでないファイルは無視
+                let path = entry.path();
+                if path.extension().unwrap_or(OsStr::new("")) != "yml" {
+                    return io::Result::Ok(ret);
+                }
+
+                // 個別のファイルの読み込みは即終了としない。
+                let read_content = self.read_file(path);
+                if read_content.is_err() {
+                    let errmsg = format!(
+                        "fail to read file: {}\n{} ",
+                        entry.path().display(),
+                        read_content.unwrap_err()
+                    );
+                    if configs::CONFIG.read().unwrap().args.is_present("verbose") {
+                        AlertMessage::warn(&mut BufWriter::new(std::io::stderr().lock()), &errmsg)?;
+                    }
+                    if !*QUIET_ERRORS_FLAG {
+                        ERROR_LOG_STACK
+                            .lock()
+                            .unwrap()
+                            .push(format!("[WARN] {}", errmsg));
+                    }
+                    self.errorrule_count += 1;
+                    return io::Result::Ok(ret);
+                }
+
+                // ここも個別のファイルの読み込みは即終了としない。
+                let yaml_contents = YamlLoader::load_from_str(&read_content.unwrap());
+                if yaml_contents.is_err() {
+                    let errmsg = format!(
+                        "Failed to parse yml: {}\n{} ",
+                        entry.path().display(),
+                        yaml_contents.unwrap_err()
+                    );
+                    if configs::CONFIG.read().unwrap().args.is_present("verbose") {
+                        AlertMessage::warn(&mut BufWriter::new(std::io::stderr().lock()), &errmsg)?;
+                    }
+                    if !*QUIET_ERRORS_FLAG {
+                        ERROR_LOG_STACK
+                            .lock()
+                            .unwrap()
+                            .push(format!("[WARN] {}", errmsg));
+                    }
+                    self.errorrule_count += 1;
+                    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()
@@ -200,6 +277,22 @@ mod tests {
     use std::path::Path;
     use yaml_rust::YamlLoader;
 
+    #[test]
+    fn test_read_file_yaml() {
+        AlertMessage::create_error_log(ERROR_LOG_PATH.to_string());
+
+        let mut yaml = yaml::ParseYaml::new();
+        let exclude_ids = RuleExclude {
+            no_use_rule: HashSet::new(),
+        };
+        let _ = &yaml.read_dir(
+            "test_files/rules/yaml/1.yml".to_string(),
+            &String::default(),
+            &exclude_ids,
+        );
+        assert_eq!(yaml.files.len(), 1);
+    }
+
     #[test]
     fn test_read_dir_yaml() {
         AlertMessage::create_error_log(ERROR_LOG_PATH.to_string());