fixed get_writeable_color and write_color_buffer arguments

2022-06-22 22:45:38 +09:00
parent 717e4dec06
commit 8a47d37c64
5 changed files with 40 additions and 29 deletions
--- a/src/afterfact.rs
+++ b/src/afterfact.rs
@@ -381,8 +381,8 @@ fn emit_csv<W: std::io::Write>(

    disp_wtr_buf.clear();
    write_color_buffer(
-        disp_wtr,
-        get_writable_color(Color::Green),
+        &disp_wtr,
+        get_writable_color(Some(Color::Green)),
        "Results Summary:",
    )
    .ok();
@@ -403,10 +403,21 @@ fn emit_csv<W: std::io::Write>(
    } else {
        (reducted_record_cnt as f64) / (all_record_cnt as f64) * 100.0
    };
-    println!("Total events: {}", all_record_cnt);
+    write_color_buffer(
+        &disp_wtr,
+        get_writable_color(None),
+        &format!("Total events: {}", all_record_cnt),
+    )
+    .ok();
+    write_color_buffer(
+        &disp_wtr,
+        get_writable_color(None),
+        &format!("Data reduction: {} events ({:.2}%)",
+        reducted_record_cnt, reducted_percent),
+    )
+    .ok();
    println!(
-        "Data reduction: {} events ({:.2}%)",
-        reducted_record_cnt, reducted_percent
+        
    );
    println!();

@@ -505,7 +516,7 @@ fn _print_unique_results(

    // output total results
    write_color_buffer(
-        BufferWriter::stdout(ColorChoice::Always),
+        &BufferWriter::stdout(ColorChoice::Always),
        None,
        &format!(
            "{} {}: {}",
@@ -525,7 +536,7 @@ fn _print_unique_results(
            head_word, level_name, tail_word, counts_by_level[i]
        );
        write_color_buffer(
-            BufferWriter::stdout(ColorChoice::Always),
+            &BufferWriter::stdout(ColorChoice::Always),
            _get_output_color(color_map, level_name),
            &output_raw_str,
        )
--- a/src/detections/print.rs
+++ b/src/detections/print.rs
@@ -324,7 +324,7 @@ impl AlertMessage {
    /// ERRORメッセージを表示する関数
    pub fn alert(contents: &str) -> io::Result<()> {
        write_color_buffer(
-            BufferWriter::stderr(ColorChoice::Always),
+            &BufferWriter::stderr(ColorChoice::Always),
            None,
            &format!("[ERROR] {}", contents),
        )
@@ -333,7 +333,7 @@ impl AlertMessage {
    /// WARNメッセージを表示する関数
    pub fn warn(contents: &str) -> io::Result<()> {
        write_color_buffer(
-            BufferWriter::stderr(ColorChoice::Always),
+            &BufferWriter::stderr(ColorChoice::Always),
            None,
            &format!("[WARN] {}", contents),
        )
--- a/src/detections/utils.rs
+++ b/src/detections/utils.rs
@@ -242,7 +242,7 @@ pub fn create_rec_info(data: Value, path: String, keys: &[String]) -> EvtxRecord
 * 標準出力のカラー出力設定を指定した値に変更し画面出力を行う関数
 */
 pub fn write_color_buffer(
-    wtr: BufferWriter,
+    wtr: &BufferWriter,
    color: Option<Color>,
    output_str: &str,
 ) -> io::Result<()> {
@@ -253,11 +253,11 @@ pub fn write_color_buffer(
 }

 /// no-colorのオプションの指定があるかを確認し、指定されている場合はNoneをかえし、指定されていない場合は引数で指定されたColorをSomeでラップして返す関数
-pub fn get_writable_color(color: Color) -> Option<Color> {
+pub fn get_writable_color(color: Option<Color>) -> Option<Color> {
    if configs::CONFIG.read().unwrap().args.no_color {
        None
    } else {
-        Some(color)
+        color
    }
 }

--- a/src/main.rs
+++ b/src/main.rs
@@ -117,7 +117,7 @@ impl App {
                Ok(output) => {
                    if output != "You currently have the latest rules." {
                        write_color_buffer(
-                            BufferWriter::stdout(ColorChoice::Always),
+                            &BufferWriter::stdout(ColorChoice::Always),
                            None,
                            "Rules updated successfully.",
                        )
@@ -170,7 +170,7 @@ impl App {

        if *STATISTICS_FLAG {
            write_color_buffer(
-                BufferWriter::stdout(ColorChoice::Always),
+                &BufferWriter::stdout(ColorChoice::Always),
                None,
                "Generating Event ID Statistics",
            )
@@ -179,7 +179,7 @@ impl App {
        }
        if *LOGONSUMMARY_FLAG {
            write_color_buffer(
-                BufferWriter::stdout(ColorChoice::Always),
+                &BufferWriter::stdout(ColorChoice::Always),
                None,
                "Generating Logons Summary",
            )
@@ -262,7 +262,7 @@ impl App {
            return;
        } else {
            write_color_buffer(
-                BufferWriter::stdout(ColorChoice::Always),
+                &BufferWriter::stdout(ColorChoice::Always),
                None,
                &configs::CONFIG.read().unwrap().headless_help,
            )
@@ -274,7 +274,7 @@ impl App {
        let analysis_duration = analysis_end_time.signed_duration_since(analysis_start_time);
        println!();
        write_color_buffer(
-            BufferWriter::stdout(ColorChoice::Always),
+            &BufferWriter::stdout(ColorChoice::Always),
            None,
            &format!("Elapsed Time: {}", &analysis_duration.hhmmssxxx()),
        )
@@ -329,15 +329,15 @@ impl App {
                    )
                    .ok();
                });
-                write_color_buffer(BufferWriter::stdout(ColorChoice::Always), None, &output).ok();
+                write_color_buffer(&BufferWriter::stdout(ColorChoice::Always), None, &output).ok();
            } else {
                //標準出力の場合
                let output = "The following pivot keywords were found:".to_string();
-                write_color_buffer(BufferWriter::stdout(ColorChoice::Always), None, &output).ok();
+                write_color_buffer(&BufferWriter::stdout(ColorChoice::Always), None, &output).ok();

                pivot_key_unions.iter().for_each(|(key, pivot_keyword)| {
                    write_color_buffer(
-                        BufferWriter::stdout(ColorChoice::Always),
+                        &BufferWriter::stdout(ColorChoice::Always),
                        None,
                        &create_output(String::default(), key, pivot_keyword),
                    )
@@ -425,7 +425,7 @@ impl App {
    fn print_contributors(&self) {
        match fs::read_to_string("./contributors.txt") {
            Ok(contents) => {
-                write_color_buffer(BufferWriter::stdout(ColorChoice::Always), None, &contents).ok();
+                write_color_buffer(&BufferWriter::stdout(ColorChoice::Always), None, &contents).ok();
            }
            Err(err) => {
                AlertMessage::alert(&format!("{}", err)).ok();
@@ -441,7 +441,7 @@ impl App {
            .min_level
            .to_uppercase();
        write_color_buffer(
-            BufferWriter::stdout(ColorChoice::Always),
+            &BufferWriter::stdout(ColorChoice::Always),
            None,
            &format!("Analyzing event files: {:?}", evtx_files.len()),
        )
@@ -667,7 +667,7 @@ impl App {
            Some(Color::Green)
        };
        write_color_buffer(
-            BufferWriter::stdout(ColorChoice::Always),
+            &BufferWriter::stdout(ColorChoice::Always),
            output_color,
            &content,
        )
@@ -686,7 +686,7 @@ impl App {
            None => {}
            Some(path) => {
                let content = fs::read_to_string(path).unwrap_or_default();
-                write_color_buffer(BufferWriter::stdout(ColorChoice::Always), None, &content).ok();
+                write_color_buffer(&BufferWriter::stdout(ColorChoice::Always), None, &content).ok();
            }
        }
    }
@@ -700,7 +700,7 @@ impl App {
        let hayabusa_rule_repo = Repository::open(Path::new("rules"));
        if hayabusa_repo.is_err() && hayabusa_rule_repo.is_err() {
            write_color_buffer(
-                BufferWriter::stdout(ColorChoice::Always),
+                &BufferWriter::stdout(ColorChoice::Always),
                None,
                "Attempting to git clone the hayabusa-rules repository into the rules folder.",
            )
@@ -879,7 +879,7 @@ impl App {
                .entry(tmp[3].to_string())
                .or_insert(0b0) += 1;
            write_color_buffer(
-                BufferWriter::stdout(ColorChoice::Always),
+                &BufferWriter::stdout(ColorChoice::Always),
                None,
                &format!(
                    "[Updated] {} (Modified: {} | Path: {})",
@@ -896,7 +896,7 @@ impl App {
            Ok("Rule updated".to_string())
        } else {
            write_color_buffer(
-                BufferWriter::stdout(ColorChoice::Always),
+                &BufferWriter::stdout(ColorChoice::Always),
                None,
                "You currently have the latest rules.",
            )
--- a/src/options/level_tuning.rs
+++ b/src/options/level_tuning.rs
@@ -59,7 +59,7 @@ impl LevelTuning {
        for (path, rule) in rulefile_loader.files {
            if let Some(new_level) = tuning_map.get(rule["id"].as_str().unwrap()) {
                write_color_buffer(
-                    BufferWriter::stdout(ColorChoice::Always),
+                    &BufferWriter::stdout(ColorChoice::Always),
                    None,
                    &format!("path: {}", path),
                )
@@ -94,7 +94,7 @@ impl LevelTuning {
                file.write_all(content.as_bytes()).unwrap();
                file.flush().unwrap();
                write_color_buffer(
-                    BufferWriter::stdout(ColorChoice::Always),
+                    &BufferWriter::stdout(ColorChoice::Always),
                    None,
                    &format!(
                        "level: {} -> {}",