diff --git a/src/main.rs b/src/main.rs index 6191fef32..2c4e0c5d8 100644 --- a/src/main.rs +++ b/src/main.rs @@ -24,7 +24,7 @@ use hayabusa::detections::utils::{ check_setting_path, get_writable_color, output_and_data_stack_for_html, output_duration, output_profile_name, }; -use hayabusa::options; +use hayabusa::{options, yaml}; use hayabusa::options::htmlreport::{self, HTML_REPORTER}; use hayabusa::options::pivot::create_output; use hayabusa::options::pivot::PIVOT_KEYWORD; @@ -50,6 +50,7 @@ use std::path::Path; use std::ptr::null_mut; use std::sync::Arc; use std::time::Duration; +use std::u128; use std::{ env, fs::{self, File}, @@ -1012,8 +1013,66 @@ impl App { || stored_static.computer_metrics_flag || stored_static.output_option.as_ref().unwrap().no_wizard) { + let mut rule_counter_wizard_map = HashMap::new(); + yaml::count_rules(&stored_static.output_option.as_ref().unwrap().rules, &filter::exclude_ids(stored_static), stored_static, &mut rule_counter_wizard_map); + let level_map:HashMap<&str, u128> = HashMap::from([ + ("INFORMATIONAL", 1), + ("LOW", 2), + ("MEDIUM", 3), + ("HIGH", 4), + ("CRITICAL", 5), + ]); println!("Scan wizard:"); println!(); + let calcurate_wizard_rule_count = |exclude_noisytarget_flag: bool, exclude_noisy_status: Vec<&str>, min_level: &str, target_status: Vec<&str>, target_tags: Vec<&str>| -> HashMap { + let mut ret = HashMap::new(); + if exclude_noisytarget_flag { + for s in exclude_noisy_status { + let mut ret_cnt = 0; + if let Some(target_status_count) = rule_counter_wizard_map.get(s) { + target_status_count.iter().for_each(| (rule_level, value)| { + let doc_level_num = level_map.get(rule_level.to_uppercase().as_str()).unwrap_or(&1); + let args_level_num = level_map.get(min_level.to_uppercase().as_str()).unwrap_or(&1); + if doc_level_num >= args_level_num { + ret_cnt += value.iter().map(|(_, cnt )| cnt).sum::() + } + }); + } + ret.insert(CompactString::from(s), ret_cnt); + } + } else { + let all_status_flag = target_status.contains(&"*"); + for s in rule_counter_wizard_map.keys() { + // 指定されたstatusに合致しないものは集計をスキップする + if (exclude_noisy_status.contains(&s.as_str()) || !target_status.contains(&s.as_str())) && !all_status_flag { + continue; + } + let mut ret_cnt = 0; + if let Some(target_status_count) = rule_counter_wizard_map.get(s) { + target_status_count.iter().for_each(| (rule_level, value)| { + let doc_level_num = level_map.get(rule_level.to_uppercase().as_str()).unwrap_or(&1); + let args_level_num = level_map.get(min_level.to_uppercase().as_str()).unwrap_or(&1); + if doc_level_num >= args_level_num { + if !target_tags.is_empty() { + for (tag, cnt) in value.iter() { + if target_tags.contains(&tag.as_str()) { + let matched_tag_cnt = ret.entry(tag.clone()); + *matched_tag_cnt.or_insert(0) += cnt; + } + } + } else { + ret_cnt += value.iter().map(|(_, cnt )| cnt).sum::() + } + } + }); + if !exclude_noisy_status.contains(&s.as_str()) { + ret.insert(s.clone(), ret_cnt); + } + } + } + } + ret + }; let selections_status = &[ ("1. Core ( status: test, stable | level: high, critical )", (vec!["test", "stable"], "high")), ("2. Core+ ( status: test, stable | level: medium, high, critical )", (vec!["test", "stable"], "medium")), @@ -1022,23 +1081,35 @@ impl App { ("5. All event and alert rules ( status: * | level: informational+ )", (vec!["*"], "informational")), ]; - let selections = selections_status.iter().map(|x| x.0).collect_vec(); - let selected_index = Select::with_theme(&ColorfulTheme::default()) - .with_prompt("Which set of detection rules would you like to load?") - .default(0) - .items(selections.as_slice()) - .interact() - .unwrap(); - status_append_output = Some(format!( - "- selected detection rule sets: {}", - selections_status[selected_index].0 - )); - stored_static.output_option.as_mut().unwrap().min_level = - selections_status[selected_index].1 .1.into(); + let sections_rule_cnt = selections_status.iter().map(|(_, (status, min_level))| { + calcurate_wizard_rule_count(false, [].to_vec(), min_level, status.to_vec(), [].to_vec()) + }).collect_vec(); + let selection_status_items = &[ + format!("1. Core ({} rules) ( status: test, stable | level: high, critical )", sections_rule_cnt[0].iter().map(|(_, cnt)| cnt).sum::()), + format!("2. Core+ ({} rules) ( status: test, stable | level: medium, high, critical )", sections_rule_cnt[1].iter().map(|(_, cnt)| cnt).sum::()), + format!("3. Core++ ({} rules) ( status: experimental, test, stable | level: medium, high, critical )", sections_rule_cnt[2].iter().map(|(_, cnt)| cnt).sum::()), + format!("4. All alert rules ({} rules) ( status: * | level: low+ )", sections_rule_cnt[3].iter().map(|(_, cnt)| cnt).sum::()), + format!("5. All event and alert rules ({} rules) ( status: * | level: informational+ )", sections_rule_cnt[4].iter().map(|(_, cnt)| cnt).sum::()) + ]; - stored_static.include_status.extend( - selections_status[selected_index] - .1 + let selected_index = Select::with_theme(&ColorfulTheme::default()) + .with_prompt("Which set of detection rules would you like to load?") + .default(0) + .items(selection_status_items.as_slice()) + .interact() + .unwrap(); + status_append_output = Some(format!( + "- selected detection rule sets: {}", + selections_status[selected_index].0 + )); + stored_static.output_option.as_mut().unwrap().min_level = + selections_status[selected_index].1 .1.into(); + + let exclude_noisy_cnt = calcurate_wizard_rule_count(true, ["exclude", "noisy", "deprecated", "unsupported"].to_vec(), selections_status[selected_index].1.1, [].to_vec(), [].to_vec()); + + stored_static.include_status.extend( + selections_status[selected_index] + .1 .0 .iter() .map(|x| x.to_owned().into()), @@ -1046,10 +1117,12 @@ impl App { let mut output_option = stored_static.output_option.clone().unwrap(); let exclude_tags = output_option.exclude_tag.get_or_insert_with(Vec::new); + let tags_cnt = calcurate_wizard_rule_count(false, [].to_vec(), selections_status[selected_index].1.1, selections_status[selected_index].1.0.clone(), ["detection.emerging_threats", "detection.threat_hunting", "sysmon"].to_vec()); // If anything other than "4. All alert rules" or "5. All event and alert rules" was selected, ask questions about tags. if selected_index < 3 { + let prompt_fmt = format!("Include Emerging Threats rules? ({} rules)", tags_cnt.get("detection.emerging_threats").unwrap_or(&0)); let et_rules_load_flag = Confirm::with_theme(&ColorfulTheme::default()) - .with_prompt("Include Emerging Threats rules?") + .with_prompt(prompt_fmt) .default(true) .show_default(true) .interact() @@ -1058,8 +1131,9 @@ impl App { if !et_rules_load_flag { exclude_tags.push("detection.emerging_threats".into()); } + let prompt_fmt = format!("Include Threat Hunting rules? ({} rules)", tags_cnt.get("detection.threat_hunting").unwrap_or(&0)); let th_rules_load_flag = Confirm::with_theme(&ColorfulTheme::default()) - .with_prompt("Include Threat Hunting rules?") + .with_prompt(prompt_fmt) .default(false) .show_default(true) .interact() @@ -1070,8 +1144,9 @@ impl App { } } // deprecated rules load prompt + let prompt_fmt = format!("Include deprecated rules? ({} rules)", exclude_noisy_cnt.get("deprecated").unwrap_or(&0)); let dep_rules_load_flag = Confirm::with_theme(&ColorfulTheme::default()) - .with_prompt("Include deprecated rules?") + .with_prompt(prompt_fmt) .default(false) .show_default(true) .interact() @@ -1085,8 +1160,9 @@ impl App { } // noisy rules load prompt + let prompt_fmt = format!("Include noisy rules? ({} rules)", tags_cnt.get("noisy").unwrap_or(&0)); let noisy_rules_load_flag = Confirm::with_theme(&ColorfulTheme::default()) - .with_prompt("Include noisy rules?") + .with_prompt(prompt_fmt) .default(false) .show_default(true) .interact() @@ -1100,8 +1176,9 @@ impl App { } // unsupported rules load prompt + let prompt_fmt = format!("Include unsupported rules? ({} rules)", exclude_noisy_cnt.get("unsupported").unwrap_or(&0)); let unsupported_rules_load_flag = Confirm::with_theme(&ColorfulTheme::default()) - .with_prompt("Include unsupported rules?") + .with_prompt(prompt_fmt) .default(false) .show_default(true) .interact() @@ -1114,8 +1191,9 @@ impl App { .enable_unsupported_rules = true; } + let prompt_fmt = format!("Include sysmon rules? ({} rules)", tags_cnt.get("sysmon").unwrap_or(&0)); let sysmon_rules_load_flag = Confirm::with_theme(&ColorfulTheme::default()) - .with_prompt("Include sysmon rules?") + .with_prompt(prompt_fmt) .default(true) .show_default(true) .interact() diff --git a/src/yaml.rs b/src/yaml.rs index 328df971a..12573c360 100644 --- a/src/yaml.rs +++ b/src/yaml.rs @@ -8,10 +8,10 @@ use crate::detections::utils; use crate::filter::RuleExclude; use compact_str::CompactString; use hashbrown::{HashMap, HashSet}; +use itertools::Itertools; use std::ffi::OsStr; use std::fs; -use std::io; -use std::io::{BufReader, Read}; +use std::io::{self, BufReader, Read}; use std::path::{Path, PathBuf}; use yaml_rust::YamlLoader; @@ -54,7 +54,7 @@ impl ParseYaml { } } - pub fn read_file(&self, path: PathBuf) -> Result { + pub fn read_file(path: PathBuf) -> Result { let mut file_content = String::new(); let mut fr = fs::File::open(path) @@ -116,7 +116,7 @@ impl ParseYaml { } // 個別のファイルの読み込みは即終了としない。 - let read_content = self.read_file(path.as_ref().to_path_buf()); + let read_content = Self::read_file(path.as_ref().to_path_buf()); if read_content.is_err() { let errmsg = format!( "fail to read file: {}\n{} ", @@ -203,7 +203,7 @@ impl ParseYaml { } // 個別のファイルの読み込みは即終了としない。 - let read_content = self.read_file(path); + let read_content = Self::read_file(path); if read_content.is_err() { let errmsg = format!( "fail to read file: {}\n{} ", @@ -479,6 +479,172 @@ impl ParseYaml { } } + /// wizardへのルール数表示のためのstatus/level/tagsごとに階層化させてカウントする + pub fn count_rules>( + path: P, + exclude_ids: &RuleExclude, + stored_static: &StoredStatic, + result_container: &mut HashMap>> + ) -> HashMap>> { + let metadata = fs::metadata(path.as_ref()); + if metadata.is_err() { + return HashMap::default(); + } + let mut yaml_docs = vec![]; + if metadata.unwrap().file_type().is_file() { + // 拡張子がymlでないファイルは無視 + if path + .as_ref() + .to_path_buf() + .extension() + .unwrap_or_else(|| OsStr::new("")) + != "yml" + { + return HashMap::default(); + } + + // 個別のファイルの読み込みは即終了としない。 + let read_content = ParseYaml::read_file(path.as_ref().to_path_buf()); + if read_content.is_err() { + return HashMap::default(); + } + + // ここも個別のファイルの読み込みは即終了としない。 + let yaml_contents = YamlLoader::load_from_str(&read_content.unwrap()); + if yaml_contents.is_err() { + return HashMap::default(); + } + + yaml_docs.extend(yaml_contents.unwrap().into_iter().map(|yaml_content| { + let filepath = format!("{}", path.as_ref().to_path_buf().display()); + (filepath, yaml_content) + })); + } else { + let entries = fs::read_dir(path); + if entries.is_err() { + return HashMap::default(); + } + yaml_docs = entries.unwrap().try_fold(vec![], |mut ret, entry| { + let entry = entry?; + // フォルダは再帰的に呼び出す。 + if entry.file_type()?.is_dir() { + count_rules( + entry.path(), + exclude_ids, + stored_static, + result_container + ); + return io::Result::Ok(ret); + } + // ファイル以外は無視 + if !entry.file_type()?.is_file() { + return io::Result::Ok(ret); + } + + // 拡張子がymlでないファイルは無視 + let path = entry.path(); + if path.extension().unwrap_or_else(|| OsStr::new("")) != "yml" { + return io::Result::Ok(ret); + } + + let path_str = path.to_str().unwrap(); + // ignore if yml file in .git folder. + if utils::contains_str(path_str, "/.git/") + || utils::contains_str(path_str, "\\.git\\") + { + return io::Result::Ok(ret); + } + + // ignore if tool test yml file in hayabusa-rules. + if utils::contains_str(path_str, "rules/tools/sigmac/test_files") + || utils::contains_str(path_str, "rules\\tools\\sigmac\\test_files") + { + return io::Result::Ok(ret); + } + + // 個別のファイルの読み込みは即終了としない。 + let read_content = ParseYaml::read_file(path); + if read_content.is_err() { + return io::Result::Ok(ret); + } + + // ここも個別のファイルの読み込みは即終了としない。 + let yaml_contents = YamlLoader::load_from_str(&read_content.unwrap()); + if yaml_contents.is_err() { + let errmsg = format!( + "Failed to parse yml: {}\n{} ", + entry.path().display(), + yaml_contents.unwrap_err() + ); + if stored_static.verbose_flag { + AlertMessage::warn(&errmsg)?; + } + if !stored_static.quiet_errors_flag { + ERROR_LOG_STACK + .lock() + .unwrap() + .push(format!("[WARN] {errmsg}")); + } + return io::Result::Ok(ret); + } + + let yaml_contents = yaml_contents.unwrap().into_iter().map(|yaml_content| { + let filepath = format!("{}", entry.path().display()); + (filepath, yaml_content) + }); + ret.extend(yaml_contents); + io::Result::Ok(ret) + }).unwrap_or_default(); + } + yaml_docs.into_iter().for_each(|(_filepath, yaml_doc)| { + //除外されたルールは無視する + let empty = vec![]; + let rule_id = &yaml_doc["id"].as_str(); + let rule_tags_vec = yaml_doc["tags"].as_vec().unwrap_or(&empty); + let included_target_tag_vec = { + let target_wizard_tags = ["detection.emerging_threats", "detection.threat_hunting", "sysmon"]; + rule_tags_vec.iter().filter(|x| target_wizard_tags.contains(&x.as_str().unwrap_or_default())).filter_map(|s| s.as_str()).collect_vec() + }; + if rule_id.is_some() { + if let Some(v) = exclude_ids + .no_use_rule + .get(&rule_id.unwrap_or(&String::default()).to_string()) + { + let entry_key = if utils::contains_str(v, "exclude_rule") { + "excluded" + } else { + "noisy" + }; + let counter = result_container.entry(entry_key.into()).or_insert(HashMap::new()); + *counter.entry(yaml_doc["level"] + .as_str() + .unwrap_or("informational") + .to_uppercase().into()).or_insert(HashMap::new()).entry(yaml_doc["status"].as_str().unwrap_or("undefined").to_lowercase().into()).or_insert(0) += 1; + return; + } + } + + if let Some(s) = yaml_doc["status"].as_str() { + // wizard用の初期カウンティングではstatusとlevelの内容を確認したうえで以降の処理は行わないようにする + let counter = result_container.entry(s.into()).or_insert(HashMap::new()); + if included_target_tag_vec.is_empty() { + *counter.entry(yaml_doc["level"] + .as_str() + .unwrap_or("informational") + .to_uppercase().into()).or_insert(HashMap::new()).entry("other".into()).or_insert(0) += 1; + } else { + for tag in included_target_tag_vec { + *counter.entry(yaml_doc["level"] + .as_str() + .unwrap_or("informational") + .to_uppercase().into()).or_insert(HashMap::new()).entry(tag.into()).or_insert(0) += 1; + } + } + } + }); + result_container.to_owned() + } + #[cfg(test)] mod tests { @@ -492,6 +658,7 @@ mod tests { use crate::detections::configs::StoredStatic; use crate::filter; use crate::yaml; + use crate::yaml::ParseYaml; use crate::yaml::RuleExclude; use compact_str::CompactString; use hashbrown::HashMap; @@ -600,9 +767,8 @@ mod tests { #[test] fn test_read_yaml() { - let yaml = yaml::ParseYaml::new(&create_dummy_stored_static()); let path = Path::new("test_files/rules/yaml/1.yml"); - let ret = yaml.read_file(path.to_path_buf()).unwrap(); + let ret = ParseYaml::read_file(path.to_path_buf()).unwrap(); let rule = YamlLoader::load_from_str(&ret).unwrap(); for i in rule { if i["title"].as_str().unwrap() == "Sysmon Check command lines" { @@ -617,9 +783,8 @@ mod tests { #[test] fn test_failed_read_yaml() { - let yaml = yaml::ParseYaml::new(&create_dummy_stored_static()); let path = Path::new("test_files/rules/yaml/error.yml"); - let ret = yaml.read_file(path.to_path_buf()).unwrap(); + let ret = ParseYaml::read_file(path.to_path_buf()).unwrap(); let rule = YamlLoader::load_from_str(&ret); assert!(rule.is_err()); }