Toml読み込み機能実装

Cargo.lock

@@ -1071,6 +1071,15 @@ dependencies = [
  "winapi",
 ]
 
+[[package]]
+name = "toml"
+version = "0.5.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ffc92d160b1eef40665be3a05630d003936a3bc7da7421277846c2613e92c71a"
+dependencies = [
+ "serde",
+]
+
 [[package]]
 name = "unicode-width"
 version = "0.1.8"
@@ -1185,5 +1194,7 @@ dependencies = [
  "quick-xml 0.17.2",
  "regex",
  "serde",
+ "serde_derive",
  "serde_json",
+ "toml",
 ]

Cargo.toml

@@ -11,8 +11,10 @@ evtx = { git = "https://github.com/omerbenamram/evtx.git" }
 quick-xml = {version = "0.17", features = ["serialize"] }
 serde = { version = "1.0", features = ["derive"] }
 serde_json = { version = "1.0"}
+serde_derive = "1.0"
 clap = "*"
 regex = "1"
 csv = "1.1"
 base64 = "*"
 flate2 = "1.0"
+toml = "0.5"

rules/test.toml

@@ -0,0 +1,2 @@
+[rule]
+severity = "high"

src/lib.rs

@@ -1,2 +1,3 @@
 pub mod detections;
 pub mod models;
+pub mod toml;

src/main.rs

@@ -6,6 +6,7 @@ use evtx::EvtxParser;
 use quick_xml::de::DeError;
 use std::{path::PathBuf, process};
 use yamato_event_analyzer::detections::detection;
+use yamato_event_analyzer::toml;
 
 fn build_app() -> clap::App<'static, 'static> {
     let program = std::env::args()
@@ -58,3 +59,4 @@ fn parse_file(filepath: &str) {
     let mut detection = detection::Detection::new();
     &detection.start(parser);
 }
+

src/models/mod.rs

@@ -1 +1,2 @@
 pub mod event;
+pub mod rule;

src/models/rule.rs

@@ -0,0 +1,14 @@
+extern crate serde;
+use serde::Deserialize;
+
+#[derive(Debug, Deserialize)]
+pub struct Rule {
+    pub severity: Option<String>,
+    
+}
+
+#[derive(Debug, Deserialize)]
+pub struct Toml {
+    pub rule: Rule,
+    
+}

src/toml.rs

@@ -0,0 +1,76 @@
+extern crate serde_derive;
+extern crate toml;
+
+use std::fs;
+use std::io;
+use std::io::{BufReader, Read};
+use std::path::{Path, PathBuf};
+use crate::models::rule;
+
+
+pub struct ParseToml {
+    pub rules: Vec<Result<rule::Toml, toml::de::Error>>,
+}
+
+impl ParseToml {
+
+    pub fn new() -> ParseToml {
+        ParseToml {
+            rules: Vec::new(),
+        }
+    }
+
+    fn read_file(&self, path: PathBuf) -> Result<String, String> {
+        let mut file_content = String::new();
+
+        let mut fr = fs::File::open(path)
+            .map(|f| BufReader::new(f))
+            .map_err(|e| e.to_string())?;
+
+        fr.read_to_string(&mut file_content)
+            .map_err(|e| e.to_string())?;
+
+        Ok(file_content)
+    }
+
+    fn read_dir<P: AsRef<Path>>(&mut self, path: P) -> io::Result<String> {
+        Ok(fs::read_dir(path)?
+            .filter_map(|entry| {
+                let entry = entry.ok()?;
+                if entry.file_type().ok()?.is_file() {
+                    match self.read_file(entry.path()) {
+                        Ok(s) => &self.rules.push(toml::from_str(&s)),
+                        Err(e) => panic!("fail to read file: {}", e),
+                    };
+                }
+                Some("")
+            })
+            .collect())
+    }
+
+}
+
+#[cfg(test)]
+mod tests {
+
+    use crate::toml;
+
+    #[test]
+    fn test_read_toml() {
+        let mut toml = toml::ParseToml::new();
+        &toml.read_dir("rules".to_string());
+
+        for rule in toml.rules {
+            match rule {
+                Ok(_rule) => {
+                    if let Some(severity) = _rule.rule.severity {
+                        assert_eq!("high", severity);
+                    }
+                },
+                Err(_) => (),
+            }
+        }
+        
+    }
+}
+