Merge pull request #1269 from Yamato-Security/1256-no-color-not-worki…

…ng-for-progress-bar-and-scan-wizard Fixed progress bar and wizard colored output when `--no-color` option is used
Yamato-Security · Feb 4, 2024 · bda68fb · bda68fb
2 parents 15e0296 + 9c98cfc
commit bda68fb
Show file tree

Hide file tree

Showing 6 changed files with 57 additions and 19 deletions.
diff --git a/CHANGELOG-Japanese.md b/CHANGELOG-Japanese.md
@@ -19,6 +19,7 @@
 **バグ修正:**
 
 - `search`コマンドの出力に入っている不要な改行文字を削除した。 (#1253) (@hitenkoku)
+- `no-color`オプション使用時のプログレスバーとウィザードのカラー出力を修正した。 (#1256) (@hitenkoku)
 
 **その他:**
 

diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -19,6 +19,7 @@
 **Bug Fixes:**
 
 - Removed newline characters in `search` command output. (#1253) (@hitenkoku)
+- Fixed progress bar and wizard colored output when `--no-color` option is used. (#1256) (@hitenkoku)
 
 **Other:**
 

diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -54,6 +54,7 @@ num = "0.4.0"
 indexmap = "2.*"
 dialoguer = "*"
 wildmatch = "2.*"
+console = "0.15.7"
 
 [profile.dev]
 debug = 0

diff --git a/src/main.rs b/src/main.rs
@@ -8,6 +8,7 @@ use bytesize::ByteSize;
 use chrono::{DateTime, Datelike, Local, NaiveDateTime, Utc};
 use clap::Command;
 use compact_str::CompactString;
+use console::{style, Style};
 use dialoguer::Confirm;
 use dialoguer::{theme::ColorfulTheme, Select};
 use evtx::{EvtxParser, ParserSettings, RecordAllocation};
@@ -1132,7 +1133,31 @@ impl App {
                 format!("5. All event and alert rules ({} rules) ( status: * | level: informational+ )", sections_rule_cnt[4].iter().map(|(_, cnt)| cnt).sum::<i128>() - sections_rule_cnt[4].get("excluded").unwrap_or(&0))
             ];
 
-            let selected_index = Select::with_theme(&ColorfulTheme::default())
+            let color_theme = if stored_static.common_options.no_color {
+                ColorfulTheme {
+                    defaults_style: Style::new().for_stderr(),
+                    prompt_style: Style::new().for_stderr().bold(),
+                    prompt_prefix: style("?".to_string()).for_stderr(),
+                    prompt_suffix: style("›".to_string()).for_stderr(),
+                    success_prefix: style("✔".to_string()).for_stderr(),
+                    success_suffix: style("·".to_string()).for_stderr(),
+                    error_prefix: style("✘".to_string()).for_stderr(),
+                    error_style: Style::new().for_stderr(),
+                    hint_style: Style::new().for_stderr(),
+                    values_style: Style::new().for_stderr(),
+                    active_item_style: Style::new().for_stderr(),
+                    inactive_item_style: Style::new().for_stderr(),
+                    active_item_prefix: style("❯".to_string()).for_stderr(),
+                    inactive_item_prefix: style(" ".to_string()).for_stderr(),
+                    checked_item_prefix: style("✔".to_string()).for_stderr(),
+                    unchecked_item_prefix: style("⬚".to_string()).for_stderr(),
+                    picked_item_prefix: style("❯".to_string()).for_stderr(),
+                    unpicked_item_prefix: style(" ".to_string()).for_stderr(),
+                }
+            } else {
+                ColorfulTheme::default()
+            };
+            let selected_index = Select::with_theme(&color_theme)
                 .with_prompt("Which set of detection rules would you like to load?")
                 .default(0)
                 .items(selection_status_items.as_slice())
@@ -1179,7 +1204,7 @@ impl App {
             if selected_index < 3 {
                 if let Some(et_cnt) = tags_cnt.get("detection.emerging_threats") {
                     let prompt_fmt = format!("Include Emerging Threats rules? ({} rules)", et_cnt);
-                    let et_rules_load_flag = Confirm::with_theme(&ColorfulTheme::default())
+                    let et_rules_load_flag = Confirm::with_theme(&color_theme)
                         .with_prompt(prompt_fmt)
                         .default(true)
                         .show_default(true)
@@ -1192,7 +1217,7 @@ impl App {
                 }
                 if let Some(th_cnt) = tags_cnt.get("detection.threat_hunting") {
                     let prompt_fmt = format!("Include Threat Hunting rules? ({} rules)", th_cnt);
-                    let th_rules_load_flag = Confirm::with_theme(&ColorfulTheme::default())
+                    let th_rules_load_flag = Confirm::with_theme(&color_theme)
                         .with_prompt(prompt_fmt)
                         .default(false)
                         .show_default(true)
@@ -1208,7 +1233,7 @@ impl App {
                 if let Some(dep_cnt) = exclude_noisy_cnt.get("deprecated") {
                     // deprecated rules load prompt
                     let prompt_fmt = format!("Include deprecated rules? ({} rules)", dep_cnt);
-                    let dep_rules_load_flag = Confirm::with_theme(&ColorfulTheme::default())
+                    let dep_rules_load_flag = Confirm::with_theme(&color_theme)
                         .with_prompt(prompt_fmt)
                         .default(false)
                         .show_default(true)
@@ -1225,13 +1250,12 @@ impl App {
                 if let Some(unsup_cnt) = exclude_noisy_cnt.get("unsupported") {
                     // unsupported rules load prompt
                     let prompt_fmt = format!("Include unsupported rules? ({} rules)", unsup_cnt);
-                    let unsupported_rules_load_flag =
-                        Confirm::with_theme(&ColorfulTheme::default())
-                            .with_prompt(prompt_fmt)
-                            .default(false)
-                            .show_default(true)
-                            .interact()
-                            .unwrap();
+                    let unsupported_rules_load_flag = Confirm::with_theme(&color_theme)
+                        .with_prompt(prompt_fmt)
+                        .default(false)
+                        .show_default(true)
+                        .interact()
+                        .unwrap();
                     if unsupported_rules_load_flag {
                         stored_static
                             .output_option
@@ -1245,7 +1269,7 @@ impl App {
             if let Some(noisy_cnt) = exclude_noisy_cnt.get("noisy") {
                 // noisy rules load prompt
                 let prompt_fmt = format!("Include noisy rules? ({} rules)", noisy_cnt);
-                let noisy_rules_load_flag = Confirm::with_theme(&ColorfulTheme::default())
+                let noisy_rules_load_flag = Confirm::with_theme(&color_theme)
                     .with_prompt(prompt_fmt)
                     .default(false)
                     .show_default(true)
@@ -1262,7 +1286,7 @@ impl App {
 
             if let Some(sysmon_cnt) = tags_cnt.get("sysmon") {
                 let prompt_fmt = format!("Include sysmon rules? ({} rules)", sysmon_cnt);
-                let sysmon_rules_load_flag = Confirm::with_theme(&ColorfulTheme::default())
+                let sysmon_rules_load_flag = Confirm::with_theme(&color_theme)
                     .with_prompt(prompt_fmt)
                     .default(true)
                     .show_default(true)
@@ -1363,8 +1387,11 @@ impl App {
             }
         }
 
-        let template =
-            "[{elapsed_precise}] {human_pos} / {human_len} {spinner:.green} [{bar:40.green}] {percent}%\r\n\r\n{msg}";
+        let template = if stored_static.common_options.no_color {
+            "[{elapsed_precise}] {human_pos} / {human_len} {spinner} [{bar:40}] {percent}%\r\n\r\n{msg}"
+        } else {
+            "[{elapsed_precise}] {human_pos} / {human_len} {spinner:.green} [{bar:40.green}] {percent}%\r\n\r\n{msg}"
+        };
         let progress_style = ProgressStyle::with_template(template)
             .unwrap()
             .progress_chars("=> ");

diff --git a/src/timeline/search.rs b/src/timeline/search.rs
@@ -419,6 +419,13 @@ pub fn search_result_dsp_msg(
                 .as_str(),
             &stored_static.disp_abbr_general_values,
         );
+        let get_char_color = |output_char_color: Option<Color>| {
+            if stored_static.common_options.no_color {
+                None
+            } else {
+                output_char_color
+            }
+        };
 
         let fmted_all_field_info = all_field_info.split_whitespace().join(" ");
         let all_field_info = if output.is_some() && stored_static.multiline_flag {
@@ -494,14 +501,14 @@ pub fn search_result_dsp_msg(
                             fields.split(':').map(|x| x.split_whitespace().join(" "));
                         write_color_buffer(
                             disp_wtr.as_mut().unwrap(),
-                            Some(Color::Rgb(255, 158, 61)),
+                            get_char_color(Some(Color::Rgb(255, 158, 61))),
                             &format!("{}: ", separated_fields_data.next().unwrap()),
                             newline_flag,
                         )
                         .ok();
                         write_color_buffer(
                             disp_wtr.as_mut().unwrap(),
-                            Some(Color::Rgb(0, 255, 255)),
+                            get_char_color(Some(Color::Rgb(0, 255, 255))),
                             separated_fields_data.join(":").trim(),
                             newline_flag,
                         )
@@ -520,7 +527,7 @@ pub fn search_result_dsp_msg(
                     //タイムスタンプとイベントタイトルは同じ色で表示
                     write_color_buffer(
                         disp_wtr.as_mut().unwrap(),
-                        Some(Color::Rgb(0, 255, 0)),
+                        get_char_color(Some(Color::Rgb(0, 255, 0))),
                         record_field_data,
                         newline_flag,
                     )
@@ -538,7 +545,7 @@ pub fn search_result_dsp_msg(
                 if !newline_flag {
                     write_color_buffer(
                         disp_wtr.as_mut().unwrap(),
-                        Some(Color::Rgb(238, 102, 97)),
+                        get_char_color(Some(Color::Rgb(238, 102, 97))),
                         " · ",
                         false,
                     )