Skip to content

Commit

Permalink
fix: output correlation rule count
Browse files Browse the repository at this point in the history
  • Loading branch information
@fukusuket
fukusuket committed Jul 22, 2024
1 parent 25c1200 commit 69c60af
Show file tree
Hide file tree
Showing 2 changed files with 114 additions and 47 deletions.
37 changes: 35 additions & 2 deletions src/detections/detection.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -153,6 +153,8 @@ impl Detection {
&rulefile_loader.rulecounter,
&rulefile_loader.rule_load_cnt,
&rulefile_loader.rule_status_cnt,
&rulefile_loader.rule_cor_cnt,
&rulefile_loader.rule_cor_ref_cnt,
&parseerror_count,
stored_static,
);
Expand Down Expand Up @@ -1104,6 +1106,8 @@ impl Detection {
rc: &HashMap<CompactString, u128>,
ld_rc: &HashMap<CompactString, u128>,
st_rc: &HashMap<CompactString, u128>,
cor_rc: &HashMap<CompactString, u128>,
cor_ref_rc: &HashMap<CompactString, u128>,
err_rc: &u128,
stored_static: &StoredStatic,
) {
Expand Down Expand Up @@ -1146,8 +1150,9 @@ impl Detection {
)
.ok();
}
println!();

if !ld_rc.is_empty() {
println!();
}
let mut sorted_st_rc: Vec<(&CompactString, &u128)> = st_rc.iter().collect();
let output_opt = stored_static.output_option.as_ref().unwrap();
let enable_deprecated_flag = output_opt.enable_deprecated_rules;
Expand Down Expand Up @@ -1192,6 +1197,34 @@ impl Detection {
});
println!();

let cor_total: u128 = cor_rc.values().sum();
let cor_ref_total: u128 = cor_ref_rc.values().sum();
if cor_total != 0 {
let col = format!(
"Correlation rules: {} ({:.2}%)",
cor_total.to_formatted_string(&Locale::en),
(cor_total as f64) / (total_loaded_rule_cnt as f64) * 100.0
);
write_color_buffer(&BufferWriter::stdout(ColorChoice::Always), None, &col, true).ok();
let col_ref = format!(
"Correlation referenced rules: {} ({:.2}%)",
cor_ref_total.to_formatted_string(&Locale::en),
(cor_ref_total as f64) / (total_loaded_rule_cnt as f64) * 100.0
);
write_color_buffer(
&BufferWriter::stdout(ColorChoice::Always),
None,
&col_ref,
true,
)
.ok();
if stored_static.html_report_flag {
html_report_stock.push(format!("- {col}"));
html_report_stock.push(format!("- {col_ref}"));
}
println!();
}

let mut sorted_rc: Vec<(&CompactString, &u128)> = rc.iter().collect();
sorted_rc.sort_by(|a, b| a.0.cmp(b.0));
sorted_rc.into_iter().for_each(|(key, value)| {
Expand Down
124 changes: 79 additions & 45 deletions src/yaml.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,13 +13,15 @@ use std::ffi::OsStr;
use std::fs;
use std::io::{self, BufReader, Read};
use std::path::{Path, PathBuf};
use yaml_rust::YamlLoader;
use yaml_rust::{Yaml, YamlLoader};

pub struct ParseYaml {
pub files: Vec<(String, yaml_rust::Yaml)>,
pub rulecounter: HashMap<CompactString, u128>,
pub rule_load_cnt: HashMap<CompactString, u128>,
pub rule_status_cnt: HashMap<CompactString, u128>,
pub rule_cor_cnt: HashMap<CompactString, u128>,
pub rule_cor_ref_cnt: HashMap<CompactString, u128>,
pub errorrule_count: u128,
pub exclude_status: HashSet<String>,
pub level_map: HashMap<String, u128>,
Expand All @@ -41,6 +43,8 @@ impl ParseYaml {
("deprecated".into(), 0_u128),
("unsupported".into(), 0_u128),
]),
rule_cor_cnt: Default::default(),
rule_cor_ref_cnt: Default::default(),
errorrule_count: 0,
exclude_status: configs::convert_option_vecs_to_hs(exclude_status_vec.as_ref()),
level_map: HashMap::from([
Expand All @@ -67,6 +71,31 @@ impl ParseYaml {
Ok(file_content)
}

fn update_correlation_counts(&mut self, yaml_docs: &Vec<Yaml>) {
for doc in yaml_docs {
if let Some(correlation) = doc["correlation"].as_hash() {
let entry = self
.rule_cor_cnt
.entry(CompactString::from("correlation"))
.or_insert(0);
*entry += 1;
if let Some(rules) = correlation.get(&Yaml::String("rules".to_string())) {
if let Some(rules_list) = rules.as_vec() {
for rule in rules_list {
if let Some(rule_str) = rule.as_str() {
// Update rules count, storing each unique rule
let rule_entry = self
.rule_cor_ref_cnt
.entry(CompactString::from(rule_str))
.or_insert(0);
*rule_entry += 1;
}
}
}
}
}
}
}
pub fn read_dir<P: AsRef<Path>>(
&mut self,
path: P,
Expand Down Expand Up @@ -137,30 +166,32 @@ impl ParseYaml {
}

// ここも個別のファイルの読み込みは即終了としない。
let yaml_contents = YamlLoader::load_from_str(&read_content.unwrap());
if yaml_contents.is_err() {
let errmsg = format!(
"Failed to parse yml: {}\n{} ",
path.as_ref().to_path_buf().display(),
yaml_contents.unwrap_err()
);
if stored_static.verbose_flag {
AlertMessage::warn(&errmsg)?;
match YamlLoader::load_from_str(&read_content.unwrap()) {
Ok(contents) => {
Self::update_correlation_counts(self, &contents);
yaml_docs.extend(contents.into_iter().map(|yaml_content| {
let filepath = format!("{}", path.as_ref().to_path_buf().display());
(filepath, yaml_content)
}));
}
if !stored_static.quiet_errors_flag {
ERROR_LOG_STACK
.lock()
.unwrap()
.push(format!("[WARN] {errmsg}"));
Err(error) => {
let errmsg = format!(
"Failed to parse yml: {}\n{} ",
path.as_ref().to_path_buf().display(),
error
);
if stored_static.verbose_flag {
AlertMessage::warn(&errmsg)?;
}
if !stored_static.quiet_errors_flag {
ERROR_LOG_STACK
.lock()
.unwrap()
.push(format!("[WARN] {errmsg}"));
}
self.errorrule_count += 1;
}
self.errorrule_count += 1;
return io::Result::Ok(String::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 mut entries = fs::read_dir(path)?;
yaml_docs = entries.try_fold(vec![], |mut ret, entry| {
Expand Down Expand Up @@ -224,32 +255,35 @@ impl ParseYaml {
}

// ここも個別のファイルの読み込みは即終了としない。
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)?;
match YamlLoader::load_from_str(&read_content.unwrap()) {
Ok(contents) => {
Self::update_correlation_counts(self, &contents);
let pair = contents.into_iter().map(|yaml_content| {
let filepath = format!("{}", entry.path().display());
(filepath, yaml_content)
});
ret.extend(pair);
io::Result::Ok(ret)
}
if !stored_static.quiet_errors_flag {
ERROR_LOG_STACK
.lock()
.unwrap()
.push(format!("[WARN] {errmsg}"));
Err(error) => {
let errmsg = format!(
"Failed to parse yml: {}\n{} ",
entry.path().display(),
error
);
if stored_static.verbose_flag {
AlertMessage::warn(&errmsg)?;
}
if !stored_static.quiet_errors_flag {
ERROR_LOG_STACK
.lock()
.unwrap()
.push(format!("[WARN] {errmsg}"));
}
self.errorrule_count += 1;
io::Result::Ok(ret)
}
self.errorrule_count += 1;
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)
})?;
}
let exist_output_opt = stored_static.output_option.is_some();
Expand Down

0 comments on commit 69c60af

Please sign in to comment.