#102実装しました。 (#133)

src/detections/detection.rs

@@ -2,7 +2,7 @@ extern crate csv;
 
 use crate::detections::rule::AggResult;
 use serde_json::Value;
-use tokio::spawn;
+use tokio::{runtime::Runtime, spawn, task::JoinHandle};
 
 use crate::detections::print::MESSAGES;
 use crate::detections::rule;
@@ -10,7 +10,7 @@ use crate::detections::rule::RuleNode;
 use crate::detections::{print::AlertMessage, utils};
 use crate::yaml::ParseYaml;
 
-use std::sync::Arc;
+use std::{sync::Arc, usize};
 
 const DIRPATH_RULES: &str = "rules";
 
@@ -32,26 +32,21 @@ impl EvtxRecordInfo {
 
 // TODO テストケースかかなきゃ...
 #[derive(Debug)]
-pub struct Detection {}
+pub struct Detection {
+    rules: Vec<RuleNode>,
+}
 
 impl Detection {
-    pub fn new() -> Detection {
+    pub fn new(rules: Vec<RuleNode>) -> Detection {
-        return Detection {};
+        return Detection { rules: rules };
     }
 
-    pub fn start(&mut self, records: Vec<EvtxRecordInfo>) {
+    pub fn start(self, rt: &Runtime, records: Vec<EvtxRecordInfo>) -> Self {
-        let rules = self.parse_rule_files();
+        return rt.block_on(self.execute_rules(records));
-        if rules.is_empty() {
-            return;
-        }
-
-        let tokio_rt = utils::create_tokio_runtime();
-        tokio_rt.block_on(Detection::execute_rules(rules, records));
-        tokio_rt.shutdown_background();
     }
 
     // ルールファイルをパースします。
-    fn parse_rule_files(&self) -> Vec<RuleNode> {
+    pub fn parse_rule_files() -> Vec<RuleNode> {
         // ルールファイルのパースを実行
         let mut rulefile_loader = ParseYaml::new();
         let resutl_readdir = rulefile_loader.read_dir(DIRPATH_RULES);
@@ -81,7 +76,7 @@ impl Detection {
                 err_msgs.iter().for_each(|err_msg| {
                     AlertMessage::alert(&mut stdout, err_msg.to_string()).ok();
                 });
-                println!("");
+                println!(""); // 一行開けるためのprintln
             });
             return Option::None;
         };
@@ -95,25 +90,52 @@ impl Detection {
             .collect();
     }
 
-    async fn execute_rules(rules: Vec<RuleNode>, records: Vec<EvtxRecordInfo>) {
+    // 複数のイベントレコードに対して、複数のルールを1個実行します。
+    async fn execute_rules(mut self, records: Vec<EvtxRecordInfo>) -> Self {
         let records_arc = Arc::new(records);
-        let traiter = rules.into_iter();
+        // // 各rule毎にスレッドを作成して、スレッドを起動する。
-        // 各rule毎にスレッドを作成して、スレッドを起動する。
+        let rules = self.rules;
-        let handles = traiter.map(|rule| {
+        let handles: Vec<JoinHandle<RuleNode>> = rules
-            let records_cloned = Arc::clone(&records_arc);
+            .into_iter()
-            return spawn(async move {
+            .map(|rule| {
-                Detection::execute_rule(rule, records_cloned);
+                let records_cloned = Arc::clone(&records_arc);
-            });
+                return spawn(async move {
-        });
+                    let moved_rule = Detection::execute_rule(rule, records_cloned);
+                    return moved_rule;
+                });
+            })
+            .collect();
 
         // 全スレッドの実行完了を待機
+        let mut rules = vec![];
         for handle in handles {
-            handle.await.unwrap();
+            let ret_rule = handle.await.unwrap();
+            rules.push(ret_rule);
+        }
+
+        // この関数の先頭でrules.into_iter()を呼び出している。それにより所有権がmapのruleを経由し、execute_ruleの引数に渡しているruleに移っているので、self.rulesには所有権が無くなっている。
+        // 所有権を失ったメンバー変数を持つオブジェクトをreturnするコードを書くと、コンパイラが怒になるので(E0382という番号のコンパイルエラー)、ここでself.rulesに所有権を戻している。
+        // self.rulesが再度所有権を取り戻せるように、Detection::execute_ruleで引数に渡したruleを戻り値として返すようにしている。
+        self.rules = rules;
+
+        return self;
+    }
+
+    pub fn add_aggcondtion_msg(&self) {
+        for rule in &self.rules {
+            if !rule.has_agg_condition() {
+                continue;
+            }
+
+            let agg_results = rule.judge_satisfy_aggcondition();
+            for value in agg_results {
+                Detection::insert_agg_message(rule, value);
+            }
         }
     }
 
-    // 検知ロジックを実行します。
+    // 複数のイベントレコードに対して、ルールを1個実行します。
-    fn execute_rule(mut rule: RuleNode, records: Arc<Vec<EvtxRecordInfo>>) {
+    fn execute_rule(mut rule: RuleNode, records: Arc<Vec<EvtxRecordInfo>>) -> RuleNode {
         let records = &*records;
         let agg_condition = rule.has_agg_condition();
         for record_info in records {
@@ -124,16 +146,10 @@ impl Detection {
             // aggregation conditionが存在しない場合はそのまま出力対応を行う
             if !agg_condition {
                 Detection::insert_message(&rule, &record_info);
-                return;
             }
         }
 
-        let agg_results = rule.judge_satisfy_aggcondition();
+        return rule;
-        for value in agg_results {
-            if agg_condition {
-                Detection::insert_agg_message(&rule, value);
-            }
-        }
     }
 
     /// 条件に合致したレコードを表示するための関数

src/detections/rule/mod.rs

@@ -3,7 +3,7 @@ use crate::detections::print::Message;
 
 use chrono::{DateTime, TimeZone, Utc};
 
-use std::{collections::HashMap, sync::Arc, vec};
+use std::{collections::HashMap, fmt::Debug, sync::Arc, vec};
 
 use serde_json::Value;
 use yaml_rust::Yaml;
@@ -29,6 +29,12 @@ pub struct RuleNode {
     countdata: HashMap<String, HashMap<String, Vec<DateTime<Utc>>>>,
 }
 
+impl Debug for RuleNode {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        return Result::Ok(());
+    }
+}
+
 unsafe impl Sync for RuleNode {}
 
 impl RuleNode {

src/main.rs

@@ -1,12 +1,13 @@
 extern crate serde;
 extern crate serde_derive;
 
-use evtx::{err, EvtxParser, ParserSettings, SerializedEvtxRecord};
+use evtx::{EvtxParser, ParserSettings};
 use std::{
     fs::{self, File},
     path::PathBuf,
+    time::Instant,
+    vec,
 };
-use tokio::{spawn, task::JoinHandle};
 use yamato_event_analyzer::detections::detection;
 use yamato_event_analyzer::detections::detection::EvtxRecordInfo;
 use yamato_event_analyzer::detections::print::AlertMessage;
@@ -14,12 +15,15 @@ use yamato_event_analyzer::omikuji::Omikuji;
 use yamato_event_analyzer::{afterfact::after_fact, detections::utils};
 use yamato_event_analyzer::{detections::configs, timeline::timeline::Timeline};
 
+// 一度にtimelineやdetectionを実行する行数
+const MAX_DETECT_RECORDS: usize = 40000;
+
 fn main() {
     if let Some(filepath) = configs::CONFIG.read().unwrap().args.value_of("filepath") {
-        detect_files(vec![PathBuf::from(filepath)]);
+        analysis_files(vec![PathBuf::from(filepath)]);
     } else if let Some(directory) = configs::CONFIG.read().unwrap().args.value_of("directory") {
         let evtx_files = collect_evtxfiles(&directory);
-        detect_files(evtx_files);
+        analysis_files(evtx_files);
     } else if configs::CONFIG.read().unwrap().args.is_present("credits") {
         print_credits();
     }
@@ -69,110 +73,92 @@ fn print_credits() {
     }
 }
 
-fn detect_files(evtx_files: Vec<PathBuf>) {
+fn analysis_files(evtx_files: Vec<PathBuf>) {
-    let evnt_records = evtx_to_jsons(&evtx_files);
-
     let mut tl = Timeline::new();
-    tl.start(&evtx_files, &evnt_records);
+    let mut detection = detection::Detection::new(detection::Detection::parse_rule_files());
 
-    let mut detection = detection::Detection::new();
+    for evtx_file in evtx_files {
-    &detection.start(evnt_records);
+        let ret = analysis_file(evtx_file, tl, detection);
+        tl = ret.0;
+        detection = ret.1;
+    }
 
     after_fact();
 }
 
-// evtxファイルをjsonに変換します。
+// Windowsイベントログファイルを1ファイル分解析する。
-fn evtx_to_jsons(evtx_files: &Vec<PathBuf>) -> Vec<EvtxRecordInfo> {
+fn analysis_file(
-    // EvtxParserを生成する。
+    evtx_filepath: PathBuf,
-    let evtx_parsers: Vec<EvtxParser<File>> = evtx_files
+    mut tl: Timeline,
-        .clone()
+    mut detection: detection::Detection,
-        .into_iter()
+) -> (Timeline, detection::Detection) {
-        .filter_map(|evtx_file| {
+    let filepath_disp = evtx_filepath.display();
-            // convert to evtx parser
+    let parser = evtx_to_jsons(evtx_filepath.clone());
-            // println!("PathBuf:{}", evtx_file.display());
+    if parser.is_none() {
-            match EvtxParser::from_path(evtx_file) {
+        return (tl, detection);
-                Ok(parser) => Option::Some(parser),
-                Err(e) => {
-                    eprintln!("{}", e);
-                    return Option::None;
-                }
-            }
-        })
-        .collect();
-
-    let tokio_rt = utils::create_tokio_runtime();
-    let ret = tokio_rt.block_on(evtx_to_json(evtx_parsers, &evtx_files));
-    tokio_rt.shutdown_background();
-
-    return ret;
-}
-
-// evtxファイルからEvtxRecordInfoを生成する。
-// 戻り値は「どのイベントファイルから生成されたXMLかを示すindex」と「変換されたXML」のタプルです。
-// タプルのindexは、引数で指定されるevtx_filesのindexに対応しています。
-async fn evtx_to_json(
-    evtx_parsers: Vec<EvtxParser<File>>,
-    evtx_files: &Vec<PathBuf>,
-) -> Vec<EvtxRecordInfo> {
-    // evtx_parser.records_json()でevtxをxmlに変換するJobを作成
-    let handles: Vec<JoinHandle<Vec<err::Result<SerializedEvtxRecord<serde_json::Value>>>>> =
-        evtx_parsers
-            .into_iter()
-            .map(|mut evtx_parser| {
-                return spawn(async move {
-                    let mut parse_config = ParserSettings::default();
-                    parse_config = parse_config.separate_json_attributes(true);
-                    parse_config = parse_config.num_threads(utils::get_thread_num());
-
-                    evtx_parser = evtx_parser.with_configuration(parse_config);
-                    let values = evtx_parser.records_json_value().collect();
-                    return values;
-                });
-            })
-            .collect();
-
-    // 作成したjobを実行し(handle.awaitの部分)、スレッドの実行時にエラーが発生した場合、標準エラー出力に出しておく
-    let mut ret = vec![];
-    for (parser_idx, handle) in handles.into_iter().enumerate() {
-        let future_result = handle.await;
-        if future_result.is_err() {
-            let evtx_filepath = &evtx_files[parser_idx].display();
-            let errmsg = format!(
-                "Failed to parse event file. EventFile:{} Error:{}",
-                evtx_filepath,
-                future_result.unwrap_err()
-            );
-            AlertMessage::alert(&mut std::io::stdout().lock(), errmsg).ok();
-            continue;
-        }
-
-        future_result.unwrap().into_iter().for_each(|parse_result| {
-            ret.push((parser_idx, parse_result));
-        });
     }
 
-    return ret
+    let mut parser = parser.unwrap();
-        .into_iter()
+    let mut records = parser.records_json_value();
-        .filter_map(|(parser_idx, parse_result)| {
+    let tokio_rt = utils::create_tokio_runtime();
+    loop {
+        let mut records_per_detect = vec![];
+        while records_per_detect.len() < MAX_DETECT_RECORDS {
             // パースに失敗している場合、エラーメッセージを出力
-            if parse_result.is_err() {
+            let next_rec = records.next();
-                let evtx_filepath = &evtx_files[parser_idx].display();
+            if next_rec.is_none() {
+                break;
+            }
+
+            let record_result = next_rec.unwrap();
+            if record_result.is_err() {
+                let evtx_filepath = &filepath_disp;
                 let errmsg = format!(
                     "Failed to parse event file. EventFile:{} Error:{}",
                     evtx_filepath,
-                    parse_result.unwrap_err()
+                    record_result.unwrap_err()
                 );
                 AlertMessage::alert(&mut std::io::stdout().lock(), errmsg).ok();
-                return Option::None;
+                continue;
             }
 
-            let record_info = EvtxRecordInfo::new(
+            let record_info =
-                evtx_files[parser_idx].display().to_string(),
+                EvtxRecordInfo::new((&filepath_disp).to_string(), record_result.unwrap().data);
-                parse_result.unwrap().data,
+            records_per_detect.push(record_info);
-            );
+        }
-            return Option::Some(record_info);
+        if records_per_detect.len() == 0 {
-        })
+            break;
-        .collect();
+        }
+
+        // timeline機能の実行
+        tl.start(&records_per_detect);
+
+        // ruleファイルの検知
+        detection = detection.start(&tokio_rt, records_per_detect);
+    }
+
+    tokio_rt.shutdown_background();
+    detection.add_aggcondtion_msg();
+
+    return (tl, detection);
+}
+
+fn evtx_to_jsons(evtx_filepath: PathBuf) -> Option<EvtxParser<File>> {
+    match EvtxParser::from_path(evtx_filepath) {
+        Ok(evtx_parser) => {
+            // parserのデフォルト設定を変更
+            let mut parse_config = ParserSettings::default();
+            parse_config = parse_config.separate_json_attributes(true); // XMLのattributeをJSONに変換する時のルールを設定
+            parse_config = parse_config.num_threads(utils::get_thread_num()); // 設定しないと遅かったので、設定しておく。
+
+            let evtx_parser = evtx_parser.with_configuration(parse_config);
+            return Option::Some(evtx_parser);
+        }
+        Err(e) => {
+            eprintln!("{}", e);
+            return Option::None;
+        }
+    }
 }
 
 fn _output_with_omikuji(omikuji: Omikuji) {

src/timeline/statistics.rs

@@ -1,5 +1,3 @@
-use std::path::PathBuf;
-
 use crate::detections::{configs, detection::EvtxRecordInfo};
 
 #[derive(Debug)]
@@ -20,11 +18,7 @@ impl EventStatistics {
     // recordからEventIDを取得するには、detection::utils::get_event_value()という関数があるので、それを使うと便利かもしれません。
 
     // 現状では、この関数の戻り値として返すVec<String>を表示するコードは実装していません。
-    pub fn start(
+    pub fn start(&mut self, _records: &Vec<EvtxRecordInfo>) -> Vec<String> {
-        &mut self,
-        evtx_files: &Vec<PathBuf>,
-        records: &Vec<EvtxRecordInfo>,
-    ) -> Vec<String> {
         // 引数でstatisticsオプションが指定されている時だけ、統計情報を出力する。
         if !configs::CONFIG
             .read()

src/timeline/timeline.rs

@@ -1,5 +1,3 @@
-use std::path::PathBuf;
-
 use crate::detections::detection::EvtxRecordInfo;
 
 use super::statistics::EventStatistics;
@@ -12,8 +10,8 @@ impl Timeline {
         return Timeline {};
     }
 
-    pub fn start(&mut self, evtx_files: &Vec<PathBuf>, records: &Vec<EvtxRecordInfo>) {
+    pub fn start(&mut self, records: &Vec<EvtxRecordInfo>) {
         let mut statistic = EventStatistics::new();
-        statistic.start(evtx_files, records);
+        statistic.start(records);
     }
 }