diff --git a/src/afterfact.rs b/src/afterfact.rs index cb209ea0..76f02d2b 100644 --- a/src/afterfact.rs +++ b/src/afterfact.rs @@ -1,29 +1,29 @@ use crate::detections::configs; use crate::detections::configs::{CURRENT_EXE_PATH, TERM_SIZE}; -use crate::detections::message::{self, LEVEL_ABBR}; use crate::detections::message::AlertMessage; +use crate::detections::message::{self, LEVEL_ABBR}; use crate::detections::utils::{self, format_time}; use crate::detections::utils::{get_writable_color, write_color_buffer}; use crate::options::profile::PROFILES; use bytesize::ByteSize; use chrono::{DateTime, Local, TimeZone, Utc}; use csv::{QuoteStyle, Writer}; -use linked_hash_map::LinkedHashMap; -use std::collections::{HashMap, HashSet, BTreeMap}; use itertools::Itertools; use krapslog::{build_sparkline, build_time_markers}; use lazy_static::lazy_static; +use linked_hash_map::LinkedHashMap; use serde::Serialize; use std::cmp::min; +use std::collections::{BTreeMap, HashMap, HashSet}; use std::error::Error; use std::fmt::Debug; -use std::{fs, collections}; use std::fs::File; use std::io; use std::io::BufWriter; use std::io::Write; use std::path::Path; use std::process; +use std::{collections, fs}; use termcolor::{BufferWriter, Color, ColorChoice, ColorSpec, WriteColor}; use terminal_size::Width; @@ -184,14 +184,7 @@ fn emit_csv( let mut detect_counts_by_computer_and_level: HashMap> = HashMap::new(); - let levels = Vec::from([ - "crit", - "high", - "med ", - "low ", - "info", - "undefined", - ]); + let levels = Vec::from(["crit", "high", "med ", "low ", "info", "undefined"]); // レベル別、日ごとの集計用変数の初期化 for level_init in levels { detect_counts_by_date_and_level.insert(level_init.to_string(), HashMap::new()); @@ -223,7 +216,12 @@ fn emit_csv( } write_color_buffer( &disp_wtr, - get_writable_color(_get_output_color(&color_map, LEVEL_ABBR.get(&detect_info.level).unwrap_or(&String::default()))), + get_writable_color(_get_output_color( + &color_map, + LEVEL_ABBR + .get(&detect_info.level) + .unwrap_or(&String::default()), + )), &_get_serialized_disp_output(detect_info.ext_field.clone(), false), false, ) @@ -381,16 +379,16 @@ enum ColPos { fn _get_serialized_disp_output(mut data: LinkedHashMap, header: bool) -> String { let data_length = &data.len(); let entries = data.entries(); - let mut ret:Vec = vec![]; + let mut ret: Vec = vec![]; if header { - entries.for_each(|entry|{ + entries.for_each(|entry| { ret.push(entry.key().to_owned()); }); } else { - entries.enumerate().for_each(|(i, entry)|{ + entries.enumerate().for_each(|(i, entry)| { if i == 0 { ret.push(_format_cellpos(entry.get(), ColPos::First)) - } else if i == data_length - 1{ + } else if i == data_length - 1 { ret.push(_format_cellpos(entry.get(), ColPos::Last)) } else { ret.push(_format_cellpos(entry.get(), ColPos::Other)) @@ -499,15 +497,19 @@ fn _print_detection_summary_by_date( tmp_cnt = *cnt; } } - wtr.set_color(ColorSpec::new().set_fg(_get_output_color(color_map, level_full_map.get(level).unwrap()))) - .ok(); + wtr.set_color(ColorSpec::new().set_fg(_get_output_color( + color_map, + level_full_map.get(level).unwrap(), + ))) + .ok(); if date_str == String::default() { max_detect_str = "n/a".to_string(); } writeln!( wtr, "Date with most total {} detections: {}", - level_full_map.get(level).unwrap(), &max_detect_str + level_full_map.get(level).unwrap(), + &max_detect_str ) .ok(); } @@ -553,12 +555,16 @@ fn _print_detection_summary_by_computer( result_vec.join(", ") }; - wtr.set_color(ColorSpec::new().set_fg(_get_output_color(color_map, level_full_map.get(level).unwrap()))) - .ok(); + wtr.set_color(ColorSpec::new().set_fg(_get_output_color( + color_map, + level_full_map.get(level).unwrap(), + ))) + .ok(); writeln!( wtr, "Top 5 computers with most unique {} detections: {}", - level_full_map.get(level).unwrap(), &result_str + level_full_map.get(level).unwrap(), + &result_str ) .ok(); } @@ -585,8 +591,8 @@ mod tests { use crate::options::profile::load_profile; use chrono::{Local, TimeZone, Utc}; use linked_hash_map::LinkedHashMap; - use std::collections::HashMap; use serde_json::Value; + use std::collections::HashMap; use std::fs::File; use std::fs::{read_to_string, remove_file}; use std::io; @@ -609,7 +615,11 @@ mod tests { let test_attack = "execution/txxxx.yyy"; let test_recinfo = "record_infoinfo11"; let test_record_id = "11111"; - let output_profile: LinkedHashMap = load_profile("test_files/config/default_profile.txt", "test_files/config/profiles.txt").unwrap(); + let output_profile: LinkedHashMap = load_profile( + "test_files/config/default_profile.txt", + "test_files/config/profiles.txt", + ) + .unwrap(); { let messages = &message::MESSAGES; messages.clear(); @@ -695,7 +705,6 @@ mod tests { } }; assert!(remove_file("./test_emit_csv.csv").is_ok()); - } #[test] @@ -736,7 +745,7 @@ mod tests { + " | " + test_recinfo + "\n"; - let mut data:LinkedHashMap = LinkedHashMap::new(); + let mut data: LinkedHashMap = LinkedHashMap::new(); data.insert("Timestamp".to_owned(), format_time(&test_timestamp, false)); data.insert("Computer".to_owned(), test_computername.to_owned()); data.insert("Channel".to_owned(), test_channel.to_owned()); @@ -747,7 +756,10 @@ mod tests { data.insert("Details".to_owned(), output.to_owned()); data.insert("RecordInformation".to_owned(), test_recinfo.to_owned()); - assert_eq!(_get_serialized_disp_output(data.clone(), true), expect_header); + assert_eq!( + _get_serialized_disp_output(data.clone(), true), + expect_header + ); assert_eq!( _get_serialized_disp_output(data.clone(), false), expect_no_header diff --git a/src/detections/detection.rs b/src/detections/detection.rs index 6a90a5d5..3f587067 100644 --- a/src/detections/detection.rs +++ b/src/detections/detection.rs @@ -209,9 +209,16 @@ impl Detection { /// 条件に合致したレコードを格納するための関数 fn insert_message(rule: &RuleNode, record_info: &EvtxRecordInfo) { let profile_all_alias = if PROFILES.is_some() { - PROFILES.as_ref().unwrap().values().cloned().collect::>().join("|") - } - else{String::default()}; + PROFILES + .as_ref() + .unwrap() + .values() + .cloned() + .collect::>() + .join("|") + } else { + String::default() + }; let tag_info: Vec = match TAGS_CONFIG.is_empty() { false => rule.yaml["tags"] .as_vec() @@ -262,7 +269,7 @@ impl Detection { } else { None }; - let level= rule.yaml["level"].as_str().unwrap_or("-").to_string(); + let level = rule.yaml["level"].as_str().unwrap_or("-").to_string(); let detect_info = DetectInfo { filepath: record_info.evtx_filepath.to_string(), rulepath: (&rule.rulepath).to_owned(), diff --git a/src/detections/message.rs b/src/detections/message.rs index 3acc6315..12eb4b35 100644 --- a/src/detections/message.rs +++ b/src/detections/message.rs @@ -6,11 +6,11 @@ use crate::detections::utils::get_serde_number_to_string; use crate::detections::utils::write_color_buffer; use chrono::{DateTime, Local, TimeZone, Utc}; use dashmap::DashMap; -use linked_hash_map::LinkedHashMap; -use std::collections::HashMap; use lazy_static::lazy_static; +use linked_hash_map::LinkedHashMap; use regex::Regex; use serde_json::Value; +use std::collections::HashMap; use std::env; use std::fs::create_dir; use std::fs::File; @@ -135,9 +135,10 @@ pub fn insert_message(detect_info: DetectInfo, event_time: DateTime) { /// メッセージを設定 pub fn insert(event_record: &Value, output: String, mut detect_info: DetectInfo) { - let parsed_detail =parse_message(event_record, output).chars() - .filter(|&c| !c.is_control()) - .collect::(); + let parsed_detail = parse_message(event_record, output) + .chars() + .filter(|&c| !c.is_control()) + .collect::(); detect_info.detail = if parsed_detail.is_empty() { "-".to_string() @@ -149,38 +150,52 @@ pub fn insert(event_record: &Value, output: String, mut detect_info: DetectInfo) let time = get_event_time(event_record).unwrap_or(default_time); for (k, v) in detect_info.ext_field.clone() { let converted_reserve_info = convert_profile_reserved_info(v, detect_info.clone(), time); - detect_info.ext_field.insert(k, parse_message(event_record, converted_reserve_info)); + detect_info + .ext_field + .insert(k, parse_message(event_record, converted_reserve_info)); } insert_message(detect_info, time) } /// profileで用いられる予約語の情報を変換する関数 -fn convert_profile_reserved_info (output:String, detect_info: DetectInfo, time: DateTime) -> String { - let config_reserved_info:HashMap = HashMap::from([ - ("Timestamp".to_string(), format_time(&time,false)), +fn convert_profile_reserved_info( + output: String, + detect_info: DetectInfo, + time: DateTime, +) -> String { + let config_reserved_info: HashMap = HashMap::from([ + ("Timestamp".to_string(), format_time(&time, false)), ("Computer".to_string(), detect_info.computername), ("Channel".to_string(), detect_info.channel), ("Level".to_string(), detect_info.level), ("EventID".to_string(), detect_info.eventid), ("MitreAttack".to_string(), detect_info.tag_info), - ("RecordID".to_string(), detect_info.record_id.unwrap_or_else(|| "-".to_string())), + ( + "RecordID".to_string(), + detect_info.record_id.unwrap_or_else(|| "-".to_string()), + ), ("RuleTitle".to_string(), detect_info.alert), ("Details".to_string(), detect_info.detail), - ("RecordInformation".to_string(), detect_info.record_information.unwrap_or_else(|| "-".to_string())), + ( + "RecordInformation".to_string(), + detect_info + .record_information + .unwrap_or_else(|| "-".to_string()), + ), ("RuleFile".to_string(), detect_info.rulepath), ("EvtxFile".to_string(), detect_info.filepath), ]); let mut ret = output; - let mut convert_target:HashMap = HashMap::new(); + let mut convert_target: HashMap = HashMap::new(); for caps in ALIASREGEX.captures_iter(&ret) { let full_target_str = &caps[0]; let target_length = full_target_str.chars().count() - 2; // The meaning of 2 is two percent let target_str = full_target_str - .chars() - .skip(1) - .take(target_length) - .collect::(); - if let Some(reserved) = config_reserved_info.get(&target_str) { + .chars() + .skip(1) + .take(target_length) + .collect::(); + if let Some(reserved) = config_reserved_info.get(&target_str) { convert_target.insert(full_target_str.to_string(), reserved.to_string()); } } @@ -385,8 +400,8 @@ impl AlertMessage { mod tests { use crate::detections::message::AlertMessage; use crate::detections::message::{parse_message, MESSAGES}; - use std::collections::HashMap; use serde_json::Value; + use std::collections::HashMap; use super::{create_output_filter_config, get_default_details}; diff --git a/src/main.rs b/src/main.rs index 98a5b4a7..5452427b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -21,7 +21,7 @@ use hayabusa::detections::pivot::PivotKeyword; use hayabusa::detections::pivot::PIVOT_KEYWORD; use hayabusa::detections::rule::{get_detection_keys, RuleNode}; use hayabusa::omikuji::Omikuji; -use hayabusa::options::{level_tuning::LevelTuning, update_rules::UpdateRules, profile::PROFILES}; +use hayabusa::options::{level_tuning::LevelTuning, profile::PROFILES, update_rules::UpdateRules}; use hayabusa::{afterfact::after_fact, detections::utils}; use hayabusa::{detections::configs, timeline::timelines::Timeline}; use hayabusa::{detections::utils::write_color_buffer, filter}; diff --git a/src/options/profile.rs b/src/options/profile.rs index a11f2466..2226a43a 100644 --- a/src/options/profile.rs +++ b/src/options/profile.rs @@ -2,8 +2,8 @@ use crate::detections::configs::{self, CURRENT_EXE_PATH}; use crate::detections::message::AlertMessage; use crate::detections::utils::check_setting_path; use crate::yaml; -use linked_hash_map::LinkedHashMap; use lazy_static::lazy_static; +use linked_hash_map::LinkedHashMap; use std::fs::File; use std::io::{BufWriter, Write}; use std::path::Path; @@ -17,12 +17,9 @@ lazy_static! { ) .to_str() .unwrap(), - check_setting_path( - &CURRENT_EXE_PATH.to_path_buf(), - "config/profiles.txt" - ) - .to_str() - .unwrap() + check_setting_path(&CURRENT_EXE_PATH.to_path_buf(), "config/profiles.txt") + .to_str() + .unwrap() ); } @@ -31,8 +28,8 @@ fn read_profile_data(profile_path: &str) -> Result, String> { let yml = yaml::ParseYaml::new(); if let Ok(loaded_profile) = yml.read_file(Path::new(profile_path).to_path_buf()) { match YamlLoader::load_from_str(&loaded_profile) { - Ok(profile_yml) => Ok(profile_yml), - Err(e) => Err(format!("Parse error: {}. {}", profile_path, e)) + Ok(profile_yml) => Ok(profile_yml), + Err(e) => Err(format!("Parse error: {}. {}", profile_path, e)), } } else { Err(format!( @@ -74,18 +71,34 @@ pub fn load_profile( let mut ret: LinkedHashMap = LinkedHashMap::new(); if let Some(profile_name) = &conf.profile { if !profile_data[profile_name.as_str()].is_badvalue() { - profile_data[profile_name.as_str()].clone().as_hash().unwrap().into_iter().for_each(|(k, v)| { - ret.insert(k.as_str().unwrap().to_string(), v.as_str().unwrap().to_string()); - }); + profile_data[profile_name.as_str()] + .clone() + .as_hash() + .unwrap() + .into_iter() + .for_each(|(k, v)| { + ret.insert( + k.as_str().unwrap().to_string(), + v.as_str().unwrap().to_string(), + ); + }); Some(ret) } else { AlertMessage::alert(&format!("Invalid profile specified: {}", profile_name)).ok(); None } } else { - profile_all[0].clone().as_hash().unwrap().into_iter().for_each(|(k, v)| { - ret.insert(k.as_str().unwrap().to_string(), v.as_str().unwrap().to_string()); - }); + profile_all[0] + .clone() + .as_hash() + .unwrap() + .into_iter() + .for_each(|(k, v)| { + ret.insert( + k.as_str().unwrap().to_string(), + v.as_str().unwrap().to_string(), + ); + }); Some(ret) } } @@ -141,8 +154,8 @@ pub fn set_default_profile(default_profile_path: &str, profile_path: &str) -> Re mod tests { use linked_hash_map::LinkedHashMap; - use crate::options::profile::load_profile; use crate::detections::configs; + use crate::options::profile::load_profile; #[test] ///オプションの設定が入ると値の冪等性が担保できないためテストを逐次的に処理する @@ -165,12 +178,21 @@ mod tests { expect.insert("RecordID".to_owned(), "%RecordID%".to_owned()); expect.insert("RuleTitle".to_owned(), "%RuleTitle%".to_owned()); expect.insert("Details".to_owned(), "%Details%".to_owned()); - expect.insert("RecordInformation".to_owned(), "%RecordInformation%".to_owned()); + expect.insert( + "RecordInformation".to_owned(), + "%RecordInformation%".to_owned(), + ); expect.insert("RuleFile".to_owned(), "%RuleFile%".to_owned()); expect.insert("EvtxFile".to_owned(), "%EvtxFile%".to_owned()); expect.insert("Tags".to_owned(), "%MitreAttack%".to_owned()); - assert_eq!(Some(expect), load_profile("test_files/config/default_profile.txt", "test_files/config/profiles.txt")); + assert_eq!( + Some(expect), + load_profile( + "test_files/config/default_profile.txt", + "test_files/config/profiles.txt" + ) + ); } /// プロファイルオプションが設定されて`おり、そのオプションに該当するプロファイルが存在する場合のテスト @@ -185,20 +207,44 @@ mod tests { expect.insert("RuleTitle".to_owned(), "%RuleTitle%".to_owned()); expect.insert("Details".to_owned(), "%Details%".to_owned()); - assert_eq!(Some(expect), load_profile("test_files/config/default_profile.txt", "test_files/config/profiles.txt")); + assert_eq!( + Some(expect), + load_profile( + "test_files/config/default_profile.txt", + "test_files/config/profiles.txt" + ) + ); } /// プロファイルオプションが設定されているが、対象のオプションが存在しない場合のテスト fn test_load_profile_no_exist_profile_files() { configs::CONFIG.write().unwrap().args.profile = Some("not_exist".to_string()); - + //両方のファイルが存在しない場合 - assert_eq!(None, load_profile("test_files/config/no_exist_default_profile.txt", "test_files/config/no_exist_profiles.txt")); + assert_eq!( + None, + load_profile( + "test_files/config/no_exist_default_profile.txt", + "test_files/config/no_exist_profiles.txt" + ) + ); //デフォルトプロファイルは存在しているがprofileオプションが指定されているため読み込み失敗の場合 - assert_eq!(None, load_profile("test_files/config/profile/default_profile.txt", "test_files/config/profile/no_exist_profiles.txt")); + assert_eq!( + None, + load_profile( + "test_files/config/profile/default_profile.txt", + "test_files/config/profile/no_exist_profiles.txt" + ) + ); //オプション先のターゲットのプロファイルファイルが存在しているが、profileオプションで指定されたオプションが存在しない場合 - assert_eq!(None, load_profile("test_files/config/no_exist_default_profile.txt", "test_files/config/profiles.txt")); + assert_eq!( + None, + load_profile( + "test_files/config/no_exist_default_profile.txt", + "test_files/config/profiles.txt" + ) + ); } }