From 21884e2a3852d98b7620a0260bb37a0366e20c66 Mon Sep 17 00:00:00 2001 From: Alan Smithee Date: Sat, 20 Mar 2021 10:04:31 +0900 Subject: [PATCH] Feature/call error message struct#66 (#69) * change way to use write trait #66 * change call error message struct #66 * erase finished TODO #66 * erase comment in error message format test #66 * resolve conflict #66 --- src/afterfact.rs | 9 +++++++-- src/detections/detection.rs | 38 +++++++++++++++++++++++++------------ src/detections/print.rs | 15 ++++++++++++--- src/main.rs | 11 +++++++++-- src/yaml.rs | 17 +++++++++++++++-- 5 files changed, 69 insertions(+), 21 deletions(-) diff --git a/src/afterfact.rs b/src/afterfact.rs index 33e32510..f2474ae7 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; @@ -25,7 +26,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 +37,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); } } diff --git a/src/detections/detection.rs b/src/detections/detection.rs index 8976bfe5..b65cdc1f 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; @@ -60,7 +61,9 @@ impl Detection { match EvtxParser::from_path(evtx_file) { Ok(parser) => Option::Some(parser), Err(e) => { - eprintln!("{}", e); + let stdout = std::io::stdout(); + let mut stdout = stdout.lock(); + AlertMessage::alert(&mut stdout, format!("{}", e)).ok(); return Option::None; } } @@ -79,7 +82,9 @@ impl Detection { if json_record.is_ok() { return Option::Some(json_record.unwrap()); } else { - eprintln!("{}", json_record.unwrap_err()); + let stdout = std::io::stdout(); + let mut stdout = stdout.lock(); + AlertMessage::alert(&mut stdout, format!("{}", json_record.unwrap_err())).ok(); return Option::None; } }) @@ -87,7 +92,9 @@ impl Detection { // serialize json from json string let result_json: Result = serde_json::from_str(&json_record.data); //// https://rust-lang-nursery.github.io/rust-cookbook/encoding/complex.html if result_json.is_err() { - eprintln!("{}", result_json.unwrap_err()); + let stdout = std::io::stdout(); + let mut stdout = stdout.lock(); + AlertMessage::alert(&mut stdout, format!("{}", result_json.unwrap_err())).ok(); return Option::None; } else { return result_json.ok(); @@ -101,7 +108,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![]; } @@ -120,14 +129,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!(""); }); diff --git a/src/detections/print.rs b/src/detections/print.rs index 37dd361f..c49012cc 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)] @@ -146,14 +147,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, "{}", contents) } } #[cfg(test)] mod tests { - use crate::detections::print::Message; + use crate::detections::print::{AlertMessage, Message}; use serde_json::Value; #[test] @@ -243,4 +244,12 @@ mod tests { 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\" }]} }"; 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(); } }; }