fix: fixed twice output progress bar bug when no config folder exist #…

…1370
Yamato-Security · Jul 25, 2024 · 62e14f7 · 62e14f7
1 parent f0a53a0
commit 62e14f7
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 8 deletions.
diff --git a/src/detections/message.rs b/src/detections/message.rs
@@ -15,6 +15,7 @@ use itertools::Itertools;
 use lazy_static::lazy_static;
 use nested::Nested;
 use regex::Regex;
+use rust_embed::Embed;
 use serde_json::Value;
 use std::env;
 use std::fs::{create_dir, File};
@@ -46,6 +47,11 @@ pub struct DetectInfo {
 
 pub struct AlertMessage {}
 
+#[derive(Embed)]
+#[folder = "config"]
+#[include = "mitre_tactics.txt"]
+struct MITRETACTICS;
+
 lazy_static! {
     #[derive(Debug,PartialEq, Eq, Ord, PartialOrd)]
     pub static ref ALIASREGEX: Regex = Regex::new(r"%[a-zA-Z0-9-_\[\]]+%").unwrap();
@@ -82,13 +88,20 @@ pub fn create_output_filter_config(
     is_lower_case: bool,
 ) -> HashMap<CompactString, CompactString> {
     let mut ret: HashMap<CompactString, CompactString> = HashMap::new();
-    let read_result = match utils::read_csv(path) {
-        Ok(c) => c,
-        Err(e) => {
-            AlertMessage::alert(&e).ok();
-            return HashMap::default();
-        }
-    };
+    let read_result;
+    if path.starts_with("config/") {
+        let mitre_tactics = MITRETACTICS::get("mitre_tactics.txt").unwrap();
+        read_result =
+            utils::parse_csv(std::str::from_utf8(mitre_tactics.data.as_ref()).unwrap_or_default());
+    } else {
+        read_result = match utils::read_csv(path) {
+            Ok(c) => c,
+            Err(e) => {
+                AlertMessage::alert(&e).ok();
+                return HashMap::default();
+            }
+        };
+    }
     read_result.iter().for_each(|line| {
         let key = if is_lower_case {
             line[0].trim().to_ascii_lowercase()

diff --git a/src/main.rs b/src/main.rs
@@ -1505,7 +1505,9 @@ impl App {
                     "{:?}",
                     &evtx_file.to_str().unwrap_or_default().replace('\\', "/")
                 );
-                pb.set_message(pb_msg);
+                if !pb_msg.is_empty() {
+                    pb.set_message(pb_msg);
+                }
             }
 
             let (detection_tmp, cnt_tmp, tl_tmp, recover_cnt_tmp, mut detect_infos) =