Skip to content

Commit

Permalink
feat: added print count of rules in the scan wizard #1206
Browse files Browse the repository at this point in the history
  • Loading branch information
@hitenkoku
hitenkoku committed Nov 17, 2023
1 parent 7006a2d commit 0cd1d80
Show file tree
Hide file tree
Showing 2 changed files with 275 additions and 32 deletions.
124 changes: 101 additions & 23 deletions src/main.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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},
Expand Down Expand Up @@ -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<CompactString, u128> {
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::<u128>()
}
});
}
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::<u128>()
}
}
});
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")),
Expand All @@ -1022,34 +1081,48 @@ 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::<u128>()),
format!("2. Core+ ({} rules) ( status: test, stable | level: medium, high, critical )", sections_rule_cnt[1].iter().map(|(_, cnt)| cnt).sum::<u128>()),
format!("3. Core++ ({} rules) ( status: experimental, test, stable | level: medium, high, critical )", sections_rule_cnt[2].iter().map(|(_, cnt)| cnt).sum::<u128>()),
format!("4. All alert rules ({} rules) ( status: * | level: low+ )", sections_rule_cnt[3].iter().map(|(_, cnt)| cnt).sum::<u128>()),
format!("5. All event and alert rules ({} rules) ( status: * | level: informational+ )", sections_rule_cnt[4].iter().map(|(_, cnt)| cnt).sum::<u128>())
];

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()),
);

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()
Expand All @@ -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()
Expand All @@ -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()
Expand All @@ -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()
Expand All @@ -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()
Expand All @@ -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()
Expand Down
Loading

0 comments on commit 0cd1d80

Please sign in to comment.