diff --git a/.gitignore b/.gitignore index 7a8a82ab..d9729367 100644 --- a/.gitignore +++ b/.gitignore @@ -2,4 +2,5 @@ /samples *.test /.vscode/ -.DS_Store \ No newline at end of file +.DS_Store +test_* \ No newline at end of file diff --git a/src/afterfact.rs b/src/afterfact.rs index 33e32510..8ac49a71 100644 --- a/src/afterfact.rs +++ b/src/afterfact.rs @@ -1,5 +1,6 @@ use crate::detections::configs; use crate::detections::print; +use crate::detections::print::AlertMessage; use chrono::{DateTime, Local, TimeZone, Utc}; use serde::Serialize; use std::error::Error; @@ -11,6 +12,7 @@ use std::process; #[serde(rename_all = "PascalCase")] pub struct CsvFormat<'a> { time: &'a str, + filepath: &'a str, title: &'a str, message: &'a str, } @@ -25,7 +27,9 @@ pub fn after_fact() { match File::create(csv_path) { Ok(file) => Box::new(file), Err(err) => { - println!("Failed to open file. {}", err); + let stdout = std::io::stdout(); + let mut stdout = stdout.lock(); + AlertMessage::alert(&mut stdout, format!("Failed to open file. {}", err)).ok(); process::exit(1); } } @@ -34,7 +38,9 @@ pub fn after_fact() { }; if let Err(err) = emit_csv(&mut target) { - println!("Failed to write CSV. {}", err); + let stdout = std::io::stdout(); + let mut stdout = stdout.lock(); + AlertMessage::alert(&mut stdout, format!("Failed to write CSV. {}", err)).ok(); process::exit(1); } } @@ -47,6 +53,7 @@ fn emit_csv(writer: &mut Box) -> Result<(), Box> { for detect_info in detect_infos { wtr.serialize(CsvFormat { time: &format_time(time), + filepath: &detect_info.filepath, title: &detect_info.title, message: &detect_info.detail, })?; @@ -79,6 +86,9 @@ where fn test_emit_csv() { use serde_json::Value; use std::fs::{read_to_string, remove_file}; + let testfilepath: &str = "test.evtx"; + let test_title = "test_title"; + let output = "pokepoke"; { let mut messages = print::MESSAGES.lock().unwrap(); @@ -86,7 +96,7 @@ fn test_emit_csv() { { "Event": { "EventData": { - "CommandLine": "hoge" + "CommandRLine": "hoge" }, "System": { "TimeCreated_attributes": { @@ -97,11 +107,27 @@ fn test_emit_csv() { } "##; let event: Value = serde_json::from_str(val).unwrap(); - messages.insert(&event, "test".to_string(), "pokepoke".to_string()); + messages.insert( + testfilepath.to_string(), + &event, + test_title.to_string(), + output.to_string(), + ); } - let expect = "Time,Title,Message -1996-02-2"; + let expect_time = Utc + .datetime_from_str("1996-02-27T01:05:01Z", "%Y-%m-%dT%H:%M:%SZ") + .unwrap(); + let expect_tz = expect_time.with_timezone(&Local); + let expect = "Time,Filepath,Title,Message\n".to_string() + + &expect_tz.clone().format("%Y-%m-%dT%H:%M:%S%:z").to_string() + + "," + + testfilepath + + "," + + test_title + + "," + + output + + "\n"; let mut file: Box = Box::new(File::create("./test_emit_csv.csv".to_string()).unwrap()); @@ -110,9 +136,8 @@ fn test_emit_csv() { match read_to_string("./test_emit_csv.csv") { Err(_) => panic!("Failed to open file"), Ok(s) => { - assert_eq!(&s[0..28], expect); + assert_eq!(s, expect); } }; - assert!(remove_file("./test_emit_csv.csv").is_ok()); } diff --git a/src/detections/detection.rs b/src/detections/detection.rs index 1dec8b51..6a10c0e4 100644 --- a/src/detections/detection.rs +++ b/src/detections/detection.rs @@ -1,5 +1,6 @@ extern crate csv; +use crate::detections::print::AlertMessage; use crate::detections::print::MESSAGES; use crate::detections::rule; use crate::detections::rule::RuleNode; @@ -11,18 +12,29 @@ use serde_json::{Error, Value}; use tokio::runtime; use tokio::{spawn, task::JoinHandle}; +use std::path::PathBuf; use std::{fs::File, sync::Arc}; -use std::{path::PathBuf, time::Instant}; const DIRPATH_RULES: &str = "rules"; +#[derive(Clone, Debug)] +pub struct EvtxRecordInfo { + evtx_filepath: String, + record: Value, +} + // TODO テストケースかかなきゃ... #[derive(Debug)] -pub struct Detection {} +pub struct Detection { + parseinfos: Vec, +} impl Detection { pub fn new() -> Detection { - Detection {} + let initializer: Vec = Vec::new(); + Detection { + parseinfos: initializer, + } } pub fn start(&mut self, evtx_files: Vec) { @@ -36,7 +48,6 @@ impl Detection { } let records = self.evtx_to_jsons(evtx_files); - runtime::Runtime::new() .unwrap() .block_on(self.execute_rule(rules, records)); @@ -48,7 +59,9 @@ impl Detection { let mut rulefile_loader = ParseYaml::new(); let resutl_readdir = rulefile_loader.read_dir(DIRPATH_RULES); if resutl_readdir.is_err() { - eprintln!("{}", resutl_readdir.unwrap_err()); + let stdout = std::io::stdout(); + let mut stdout = stdout.lock(); + AlertMessage::alert(&mut stdout, format!("{}", resutl_readdir.unwrap_err())).ok(); return vec![]; } @@ -67,14 +80,19 @@ impl Detection { err_msgs_result.err().iter().for_each(|err_msgs| { // TODO 本当はファイルパスを出力したい // ParseYamlの変更が必要なので、一旦yamlのタイトルを表示。 - - // TODO エラーの出力方法を統一したい。 - // エラー出力用のクラスを作成してもいいかも - println!( - "[ERROR] Failed to parse Rule file. (Error Rule Title : {})", - rule.yaml["title"].as_str().unwrap_or("") - ); - err_msgs.iter().for_each(|err_msg| println!("{}", err_msg)); + let stdout = std::io::stdout(); + let mut stdout = stdout.lock(); + AlertMessage::alert( + &mut stdout, + format!( + "Failed to parse Rule file. (Error Rule Title : {})", + rule.yaml["title"].as_str().unwrap_or("") + ), + ) + .ok(); + err_msgs.iter().for_each(|err_msg| { + AlertMessage::alert(&mut stdout, err_msg.to_string()).ok(); + }); println!(""); }); @@ -84,7 +102,7 @@ impl Detection { } // evtxファイルをjsonに変換します。 - fn evtx_to_jsons(&mut self, evtx_files: Vec) -> Vec { + fn evtx_to_jsons(&mut self, evtx_files: Vec) -> Vec { // EvtxParserを生成する。 let evtx_parsers: Vec> = evtx_files .iter() @@ -102,18 +120,41 @@ impl Detection { let xml_records = runtime::Runtime::new() .unwrap() - .block_on(self.evtx_to_xml(evtx_parsers)); - - return runtime::Runtime::new() + .block_on(self.evtx_to_xml(evtx_parsers, &evtx_files)); + let json_records = runtime::Runtime::new() .unwrap() - .block_on(self.xml_to_json(xml_records)); + .block_on(self.xml_to_json(xml_records, &evtx_files)); + + let mut evtx_file_index = 0; + return json_records + .into_iter() + .map(|json_records_per_evtxfile| { + let evtx_filepath = evtx_files[evtx_file_index].display().to_string(); + let ret: Vec = json_records_per_evtxfile + .into_iter() + .map(|json_record| { + return EvtxRecordInfo { + evtx_filepath: String::from(&evtx_filepath), + record: json_record, + }; + }) + .collect(); + evtx_file_index = evtx_file_index + 1; + return ret; + }) + .flatten() + .collect(); } // evtxファイルからxmlを生成する。 + // ちょっと分かりにくいですが、戻り値の型はVec>ではなくて、Vec>>になっています。 + // 2次元配列にしている理由は、この後Value型(EvtxのXMLをJSONに変換したやつ)とイベントファイルのパスをEvtxRecordInfo構造体で保持するためです。 + // EvtxParser毎にSerializedEvtxRecordをグルーピングするために2次元配列にしています。 async fn evtx_to_xml( &mut self, evtx_parsers: Vec>, - ) -> Vec> { + evtx_files: &Vec, + ) -> Vec>> { // evtx_parser.records_json()でevtxをxmlに変換するJobを作成 let handles: Vec>>>> = evtx_parsers .into_iter() @@ -130,71 +171,129 @@ impl Detection { // 作成したjobを実行し(handle.awaitの部分)、スレッドの実行時にエラーが発生した場合、標準エラー出力に出しておく let mut ret = vec![]; + let mut evtx_file_index = 0; for handle in handles { let future_result = handle.await; if future_result.is_err() { - eprintln!("{}", future_result.unwrap_err()); + let evtx_filepath = &evtx_files[evtx_file_index].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; } + evtx_file_index = evtx_file_index + 1; ret.push(future_result.unwrap()); } - // xmlの変換でエラーが出た場合、標準エラー出力に出しておく + // パースに失敗しているレコードを除外して、返す。 + // SerializedEvtxRecordがどのEvtxParserからパースされたのか分かるようにするため、2次元配列のまま返す。 + let mut evtx_file_index = 0; return ret .into_iter() - .flatten() - .filter_map(|parse_result| { - if parse_result.is_err() { - eprintln!("{}", parse_result.unwrap_err()); - return Option::None; - } - - return Option::Some(parse_result.unwrap()); + .map(|parse_results| { + let ret = parse_results + .into_iter() + .filter_map(|parse_result| { + if parse_result.is_err() { + let evtx_filepath = &evtx_files[evtx_file_index].display(); + let errmsg = format!( + "Failed to parse event file. EventFile:{} Error:{}", + evtx_filepath, + parse_result.unwrap_err() + ); + AlertMessage::alert(&mut std::io::stdout().lock(), errmsg).ok(); + return Option::None; + } + return Option::Some(parse_result.unwrap()); + }) + .collect(); + evtx_file_index = evtx_file_index + 1; + return ret; }) .collect(); } // xmlからjsonに変換します。 - async fn xml_to_json(&mut self, xml_records: Vec>) -> Vec { + async fn xml_to_json( + &mut self, + xml_records: Vec>>, + evtx_files: &Vec, + ) -> Vec> { // xmlからjsonに変換するJobを作成 - let handles: Vec>> = xml_records + let handles: Vec>>> = xml_records .into_iter() - .map(|xml_record| { - return spawn(async move { - return serde_json::from_str(&xml_record.data); - }); + .map(|xml_records| { + return xml_records + .into_iter() + .map(|xml_record| { + return spawn(async move { + return serde_json::from_str(&xml_record.data); + }); + }) + .collect(); }) .collect(); // 作成したjobを実行し(handle.awaitの部分)、スレッドの実行時にエラーが発生した場合、標準エラー出力に出しておく let mut ret = vec![]; - for handle in handles { - let future_result = handle.await; - if future_result.is_err() { - eprintln!("{}", future_result.unwrap_err()); - continue; - } - - ret.push(future_result.unwrap()); - } - - // xmlの変換でエラーが出た場合、標準エラー出力に出しておく - return ret - .into_iter() - .filter_map(|parse_result| { - if parse_result.is_err() { - eprintln!("{}", parse_result.unwrap_err()); - return Option::None; + let mut evtx_file_index = 0; + for handles_per_evtxfile in handles { + let mut sub_ret = vec![]; + for handle in handles_per_evtxfile { + let future_result = handle.await; + if future_result.is_err() { + let evtx_filepath = &evtx_files[evtx_file_index].display(); + let errmsg = format!( + "Failed to serialize from event xml to json. EventFile:{} Error:{}", + evtx_filepath, + future_result.unwrap_err() + ); + AlertMessage::alert(&mut std::io::stdout().lock(), errmsg).ok(); + continue; } - return Option::Some(parse_result.unwrap()); + sub_ret.push(future_result.unwrap()); + } + ret.push(sub_ret); + evtx_file_index = evtx_file_index + 1; + } + + // JSONの変換に失敗したものを除外して、返す。 + // ValueがどのEvtxParserからパースされたのか分かるようにするため、2次元配列のまま返す。 + let mut evtx_file_index = 0; + return ret + .into_iter() + .map(|parse_results| { + let successed = parse_results + .into_iter() + .filter_map(|parse_result| { + if parse_result.is_err() { + let evtx_filepath = &evtx_files[evtx_file_index].display(); + let errmsg = format!( + "Failed to serialize from event xml to json. EventFile:{} Error:{}", + evtx_filepath, + parse_result.unwrap_err() + ); + AlertMessage::alert(&mut std::io::stdout().lock(), errmsg).ok(); + return Option::None; + } + + return Option::Some(parse_result.unwrap()); + }) + .collect(); + evtx_file_index = evtx_file_index + 1; + + return successed; }) .collect(); } // 検知ロジックを実行します。 - async fn execute_rule(&mut self, rules: Vec, records: Vec) { + async fn execute_rule(&mut self, rules: Vec, records: Vec) { // 複数スレッドで所有権を共有するため、recordsをArcでwwap let mut records_arcs = vec![]; for record_chunk in Detection::chunks(records, num_cpus::get() * 4) { @@ -213,9 +312,9 @@ impl Detection { let handle: JoinHandle> = spawn(async move { let mut ret = vec![]; - for record in records_arc_clone.iter() { + for record_info in records_arc_clone.iter() { for rule in rules_clones.iter() { - if rule.select(record) { + if rule.select(&record_info.record) { // TODO ここはtrue/falseじゃなくて、ruleとrecordのタプルをretにpushする実装に変更したい。 ret.push(true); } else { @@ -234,22 +333,25 @@ impl Detection { for record_chunk_arc in &records_arcs { let mut handles_ret_ite = handles_ite.next().unwrap().await.unwrap().into_iter(); for rule in rules_arc.iter() { - for record_arc in record_chunk_arc.iter() { - if handles_ret_ite.next().unwrap() == true { - // TODO メッセージが多いと、rule.select()よりもこの処理の方が時間かかる。 - message.insert( - record_arc, - rule.yaml["title"].as_str().unwrap_or("").to_string(), - rule.yaml["output"].as_str().unwrap_or("").to_string(), - ); + for record_info_arc in record_chunk_arc.iter() { + if handles_ret_ite.next().unwrap() == false { + continue; } + + // TODO メッセージが多いと、rule.select()よりもこの処理の方が時間かかる。 + message.insert( + record_info_arc.evtx_filepath.to_string(), + &record_info_arc.record, + rule.yaml["title"].as_str().unwrap_or("").to_string(), + rule.yaml["output"].as_str().unwrap_or("").to_string(), + ); } } } } // 配列を指定したサイズで分割する。Vector.chunksと同じ動作をするが、Vectorの関数だとinto的なことができないので自作 - fn chunks(ary: Vec, size: usize) -> Vec> { + fn chunks(ary: Vec, size: usize) -> Vec> { let arylen = ary.len(); let mut ite = ary.into_iter(); diff --git a/src/detections/print.rs b/src/detections/print.rs index 37dd361f..c6ab0cd0 100644 --- a/src/detections/print.rs +++ b/src/detections/print.rs @@ -7,6 +7,7 @@ use regex::Regex; use serde_json::Value; use std::collections::BTreeMap; use std::collections::HashMap; +use std::io::{self, Write}; use std::sync::Mutex; #[derive(Debug)] @@ -16,6 +17,7 @@ pub struct Message { #[derive(Debug, Clone)] pub struct DetectInfo { + pub filepath: String, pub title: String, pub detail: String, } @@ -33,7 +35,13 @@ impl Message { } /// メッセージを設定 - pub fn insert(&mut self, event_record: &Value, event_title: String, output: String) { + pub fn insert( + &mut self, + target_file: String, + event_record: &Value, + event_title: String, + output: String, + ) { if output.is_empty() { return; } @@ -42,6 +50,7 @@ impl Message { let default_time = Utc.ymd(1970, 1, 1).and_hms(0, 0, 0); let time = Message::get_event_time(event_record).unwrap_or(default_time); let detect_info = DetectInfo { + filepath: target_file, title: event_title, detail: message.to_string(), }; @@ -146,14 +155,14 @@ impl Message { } impl AlertMessage { - pub fn alert(contents: String) { - println!("[ERROR] {}", contents); + pub fn alert(w: &mut W, contents: String) -> io::Result<()> { + writeln!(w, "[ERROR] {}", contents) } } #[cfg(test)] mod tests { - use crate::detections::print::Message; + use crate::detections::print::{AlertMessage, Message}; use serde_json::Value; #[test] @@ -175,6 +184,7 @@ mod tests { "##; let event_record_1: Value = serde_json::from_str(json_str_1).unwrap(); message.insert( + "a".to_string(), &event_record_1, "test1".to_string(), "CommandLine1: %CommandLine%".to_string(), @@ -196,6 +206,7 @@ mod tests { "##; let event_record_2: Value = serde_json::from_str(json_str_2).unwrap(); message.insert( + "a".to_string(), &event_record_2, "test2".to_string(), "CommandLine2: %CommandLine%".to_string(), @@ -217,6 +228,7 @@ mod tests { "##; let event_record_3: Value = serde_json::from_str(json_str_3).unwrap(); message.insert( + "a".to_string(), &event_record_3, "test3".to_string(), "CommandLine3: %CommandLine%".to_string(), @@ -233,6 +245,7 @@ mod tests { "##; let event_record_4: Value = serde_json::from_str(json_str_4).unwrap(); message.insert( + "a".to_string(), &event_record_4, "test4".to_string(), "CommandLine4: %CommandLine%".to_string(), @@ -240,7 +253,15 @@ mod tests { let display = format!("{}", format_args!("{:?}", message)); println!("display::::{}", display); - let expect = "Message { map: {1970-01-01T00:00:00Z: [DetectInfo { title: \"test4\", detail: \"CommandLine4: hoge\" }], 1996-02-27T01:05:01Z: [DetectInfo { title: \"test1\", detail: \"CommandLine1: hoge\" }, DetectInfo { title: \"test2\", detail: \"CommandLine2: hoge\" }], 2000-01-21T09:06:01Z: [DetectInfo { title: \"test3\", detail: \"CommandLine3: hoge\" }]} }"; + let expect = "Message { map: {1970-01-01T00:00:00Z: [DetectInfo { filepath: \"a\", title: \"test4\", detail: \"CommandLine4: hoge\" }], 1996-02-27T01:05:01Z: [DetectInfo { filepath: \"a\", title: \"test1\", detail: \"CommandLine1: hoge\" }, DetectInfo { filepath: \"a\", title: \"test2\", detail: \"CommandLine2: hoge\" }], 2000-01-21T09:06:01Z: [DetectInfo { filepath: \"a\", title: \"test3\", detail: \"CommandLine3: hoge\" }]} }"; assert_eq!(display, expect); } + + #[test] + fn test_error_message() { + let input = "TEST!"; + let stdout = std::io::stdout(); + let mut stdout = stdout.lock(); + AlertMessage::alert(&mut stdout, input.to_string()).expect("[ERROR] TEST!"); + } } diff --git a/src/main.rs b/src/main.rs index d467018f..79f62399 100644 --- a/src/main.rs +++ b/src/main.rs @@ -5,6 +5,7 @@ use std::{fs, path::PathBuf}; use yamato_event_analyzer::afterfact::after_fact; use yamato_event_analyzer::detections::configs; use yamato_event_analyzer::detections::detection; +use yamato_event_analyzer::detections::print::AlertMessage; use yamato_event_analyzer::omikuji::Omikuji; fn main() { @@ -21,7 +22,9 @@ fn main() { fn collect_evtxfiles(dirpath: &str) -> Vec { let entries = fs::read_dir(dirpath); if entries.is_err() { - eprintln!("{}", entries.unwrap_err()); + let stdout = std::io::stdout(); + let mut stdout = stdout.lock(); + AlertMessage::alert(&mut stdout, format!("{}", entries.unwrap_err())).ok(); return vec![]; } @@ -50,9 +53,13 @@ fn collect_evtxfiles(dirpath: &str) -> Vec { } fn print_credits() { + let stdout = std::io::stdout(); + let mut stdout = stdout.lock(); match fs::read_to_string("./credits.txt") { Ok(contents) => println!("{}", contents), - Err(err) => println!("{}", err), + Err(err) => { + AlertMessage::alert(&mut stdout, format!("{}", err)).ok(); + } } } diff --git a/src/yaml.rs b/src/yaml.rs index a3b7323d..cd44e52a 100644 --- a/src/yaml.rs +++ b/src/yaml.rs @@ -1,6 +1,7 @@ extern crate serde_derive; extern crate yaml_rust; +use crate::detections::print::AlertMessage; use std::fs; use std::io; use std::io::{BufReader, Read}; @@ -34,6 +35,8 @@ impl ParseYaml { .filter_map(|entry| { let entry = entry.ok()?; if entry.file_type().ok()?.is_file() { + let stdout = std::io::stdout(); + let mut stdout = stdout.lock(); match self.read_file(entry.path()) { Ok(s) => { match YamlLoader::load_from_str(&s) { @@ -45,11 +48,21 @@ impl ParseYaml { } } } - Err(e) => eprintln!("fail to read file\n{}\n{} ", s, e), + Err(e) => { + AlertMessage::alert( + &mut stdout, + format!("fail to read file\n{}\n{} ", s, e), + ) + .ok(); + } } } Err(e) => { - eprintln!("fail to read file: {}\n{} ", entry.path().display(), e) + AlertMessage::alert( + &mut stdout, + format!("fail to read file: {}\n{} ", entry.path().display(), e), + ) + .ok(); } }; }