To speed up

This commit is contained in:
DastInDark
2022-07-25 06:14:00 +09:00
parent b22bf5d13f
commit 798214cb38
4 changed files with 178 additions and 132 deletions

View File

@@ -208,7 +208,7 @@ fn emit_csv<W: std::io::Write>(
write_color_buffer(
&disp_wtr,
get_writable_color(None),
&_get_serialized_disp_output(PROFILES.as_ref().unwrap().clone(), true),
&_get_serialized_disp_output(&mut PROFILES.as_ref().unwrap().clone(), true),
false,
)
.ok();
@@ -222,7 +222,7 @@ fn emit_csv<W: std::io::Write>(
.get(&detect_info.level)
.unwrap_or(&String::default()),
)),
&_get_serialized_disp_output(detect_info.ext_field.clone(), false),
&_get_serialized_disp_output(&mut detect_info.ext_field.clone(), false),
false,
)
.ok();
@@ -279,7 +279,7 @@ fn emit_csv<W: std::io::Write>(
wtr.flush()?;
}
let output_path = configs::CONFIG.read().unwrap().args.output.clone();
let output_path = &configs::CONFIG.read().unwrap().args.output;
if let Some(path) = output_path {
if let Ok(metadata) = fs::metadata(path) {
println!(
@@ -376,7 +376,7 @@ enum ColPos {
Other,
}
fn _get_serialized_disp_output(mut data: LinkedHashMap<String, String>, header: bool) -> String {
fn _get_serialized_disp_output(data: &mut LinkedHashMap<String, String>, header: bool) -> String {
let data_length = &data.len();
let entries = data.entries();
let mut ret: Vec<String> = vec![];
@@ -475,7 +475,6 @@ fn _print_detection_summary_by_date(
let mut wtr = buf_wtr.buffer();
wtr.set_color(ColorSpec::new().set_fg(None)).ok();
let output_levels = Vec::from(["crit", "high", "med ", "low ", "info"]);
let level_full_map = HashMap::from([
("crit", "critical"),
("high", "high"),
@@ -484,7 +483,7 @@ fn _print_detection_summary_by_date(
("info", "informational"),
]);
for level in output_levels {
for level in LEVEL_ABBR.values() {
// output_levelsはlevelsからundefinedを除外した配列であり、各要素は必ず初期化されているのでSomeであることが保証されているのでunwrapをそのまま実施
let detections_by_day = detect_counts_by_date.get(level).unwrap();
let mut max_detect_str = String::default();
@@ -499,7 +498,7 @@ fn _print_detection_summary_by_date(
}
wtr.set_color(ColorSpec::new().set_fg(_get_output_color(
color_map,
level_full_map.get(level).unwrap(),
level_full_map.get(level.as_str()).unwrap(),
)))
.ok();
if date_str == String::default() {
@@ -508,7 +507,7 @@ fn _print_detection_summary_by_date(
writeln!(
wtr,
"Date with most total {} detections: {}",
level_full_map.get(level).unwrap(),
level_full_map.get(level.as_str()).unwrap(),
&max_detect_str
)
.ok();
@@ -525,7 +524,6 @@ fn _print_detection_summary_by_computer(
let mut wtr = buf_wtr.buffer();
wtr.set_color(ColorSpec::new().set_fg(None)).ok();
let output_levels = Vec::from(["crit", "high", "med ", "low ", "info"]);
let level_full_map = HashMap::from([
("crit", "critical"),
("high", "high"),
@@ -534,7 +532,7 @@ fn _print_detection_summary_by_computer(
("info", "informational"),
]);
for level in output_levels {
for level in LEVEL_ABBR.values() {
// output_levelsはlevelsからundefinedを除外した配列であり、各要素は必ず初期化されているのでSomeであることが保証されているのでunwrapをそのまま実施
let detections_by_computer = detect_counts_by_computer.get(level).unwrap();
let mut result_vec: Vec<String> = Vec::new();
@@ -557,13 +555,13 @@ fn _print_detection_summary_by_computer(
wtr.set_color(ColorSpec::new().set_fg(_get_output_color(
color_map,
level_full_map.get(level).unwrap(),
level_full_map.get(level.as_str()).unwrap(),
)))
.ok();
writeln!(
wtr,
"Top 5 computers with most unique {} detections: {}",
level_full_map.get(level).unwrap(),
level_full_map.get(level.as_str()).unwrap(),
&result_str
)
.ok();
@@ -615,6 +613,10 @@ mod tests {
let test_attack = "execution/txxxx.yyy";
let test_recinfo = "record_infoinfo11";
let test_record_id = "11111";
let expect_time = Utc
.datetime_from_str("1996-02-27T01:05:01Z", "%Y-%m-%dT%H:%M:%SZ")
.unwrap();
let expect_tz = expect_time.with_timezone(&Local);
let output_profile: LinkedHashMap<String, String> = load_profile(
"test_files/config/default_profile.txt",
"test_files/config/profiles.txt",
@@ -638,6 +640,26 @@ mod tests {
}
"##;
let event: Value = serde_json::from_str(val).unwrap();
let mut profile_converter: HashMap<String, String> = HashMap::from([
("%Timestamp%".to_owned(), format_time(&expect_time, false)),
("%Computer%".to_owned(), test_computername.to_string()),
(
"%Channel%".to_owned(),
mock_ch_filter
.get("Security")
.unwrap_or(&String::default())
.to_string(),
),
("%Level%".to_owned(), test_level.to_string()),
("%EventID%".to_owned(), test_eventid.to_string()),
("%MitreAttack%".to_owned(), test_attack.to_string()),
("%RecordID%".to_owned(), test_record_id.to_string()),
("%RuleTitle%".to_owned(), test_title.to_owned()),
("%RecordInformation%".to_owned(), test_recinfo.to_owned()),
("%RuleFile%".to_owned(), test_rulepath.to_string()),
("%EvtxFile%".to_owned(), test_filepath.to_string()),
("%Tags%".to_owned(), test_attack.to_string()),
]);
message::insert(
&event,
output.to_string(),
@@ -658,12 +680,10 @@ mod tests {
record_id: Option::Some(test_record_id.to_string()),
ext_field: output_profile,
},
expect_time,
&mut profile_converter,
);
}
let expect_time = Utc
.datetime_from_str("1996-02-27T01:05:01Z", "%Y-%m-%dT%H:%M:%SZ")
.unwrap();
let expect_tz = expect_time.with_timezone(&Local);
let expect =
"Timestamp,Computer,Channel,Level,EventID,MitreAttack,RecordID,RuleTitle,Details,RecordInformation,RuleFile,EvtxFile,Tags\n"
.to_string()
@@ -756,12 +776,9 @@ 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(&mut data, true), expect_header);
assert_eq!(
_get_serialized_disp_output(data.clone(), true),
expect_header
);
assert_eq!(
_get_serialized_disp_output(data.clone(), false),
_get_serialized_disp_output(&mut data, false),
expect_no_header
);
}

View File

@@ -1,9 +1,13 @@
extern crate csv;
use crate::detections::configs;
use crate::detections::utils::format_time;
use crate::detections::utils::write_color_buffer;
use crate::options::profile::LOAEDED_PROFILE_ALIAS;
use crate::options::profile::PRELOAD_PROFILE;
use crate::options::profile::PRELOAD_PROFILE_REGEX;
use crate::options::profile::PROFILES;
use chrono::{TimeZone, Utc};
use termcolor::{BufferWriter, Color, ColorChoice};
use crate::detections::message::AlertMessage;
@@ -256,7 +260,78 @@ impl Detection {
} else {
None
};
let default_time = Utc.ymd(1970, 1, 1).and_hms(0, 0, 0);
let time = message::get_event_time(&record_info.record).unwrap_or(default_time);
let level = rule.yaml["level"].as_str().unwrap_or("-").to_string();
let mut profile_converter: HashMap<String, String> = HashMap::new();
for (k, v) in PROFILES.as_ref().unwrap().iter() {
let tmp = v.as_str();
for target_profile in PRELOAD_PROFILE_REGEX.matches(tmp).into_iter() {
match PRELOAD_PROFILE[target_profile] {
"%Timestamp%" => {
profile_converter.insert(k.to_string(), format_time(&time, false));
}
"%Computer%" => {
profile_converter.insert(
k.to_string(),
record_info.record["Event"]["System"]["Computer"]
.to_string()
.replace('\"', ""),
);
}
"%Channel%" => {
profile_converter.insert(
k.to_string(),
CH_CONFIG.get(ch_str).unwrap_or(ch_str).to_string(),
);
}
"%Level%" => {
profile_converter.insert(
k.to_string(),
LEVEL_ABBR.get(&level).unwrap_or(&level).to_string(),
);
}
"%EventID%" => {
profile_converter.insert(k.to_string(), eid.to_owned());
}
"%MitreAttack%" => {
profile_converter.insert(k.to_string(), tag_info.join(" | "));
}
"%RecordID%" => {
profile_converter.insert(
k.to_string(),
rec_id.as_ref().unwrap_or(&"-".to_string()).to_owned(),
);
}
"%RuleTitle%" => {
profile_converter.insert(
k.to_string(),
rule.yaml["title"].as_str().unwrap_or("").to_string(),
);
}
"%RecordInformation%" => {
profile_converter.insert(
k.to_string(),
opt_record_info
.as_ref()
.unwrap_or(&"-".to_string())
.to_owned(),
);
}
"%RuleFile%" => {
profile_converter.insert(k.to_string(), (&rule.rulepath).to_owned());
}
"%EvtxFile%" => {
profile_converter
.insert(k.to_string(), record_info.evtx_filepath.to_string());
}
_ => {}
}
}
}
let detect_info = DetectInfo {
filepath: record_info.evtx_filepath.to_string(),
rulepath: (&rule.rulepath).to_owned(),
@@ -280,6 +355,8 @@ impl Detection {
.unwrap_or(&default_output)
.to_string(),
detect_info,
time,
&mut profile_converter,
);
}

View File

@@ -4,7 +4,8 @@ use crate::detections::configs::CURRENT_EXE_PATH;
use crate::detections::utils;
use crate::detections::utils::get_serde_number_to_string;
use crate::detections::utils::write_color_buffer;
use chrono::{DateTime, Local, TimeZone, Utc};
use crate::options::profile::PROFILES;
use chrono::{DateTime, Local, Utc};
use dashmap::DashMap;
use lazy_static::lazy_static;
use linked_hash_map::LinkedHashMap;
@@ -20,8 +21,6 @@ use std::path::Path;
use std::sync::Mutex;
use termcolor::{BufferWriter, ColorChoice};
use super::utils::format_time;
#[derive(Debug, Clone)]
pub struct DetectInfo {
pub filepath: String,
@@ -83,7 +82,7 @@ lazy_static! {
.display()
));
pub static ref LEVEL_ABBR: HashMap<String, String> = HashMap::from([
(String::from("cruitical"), String::from("crit")),
(String::from("critical"), String::from("crit")),
(String::from("high"), String::from("high")),
(String::from("medium"), String::from("med ")),
(String::from("low"), String::from("low ")),
@@ -113,9 +112,8 @@ pub fn create_output_filter_config(
return;
}
let empty = &"".to_string();
let tag_full_str = line.get(0).unwrap_or(empty).trim();
let tag_replace_str = line.get(1).unwrap_or(empty).trim();
let tag_full_str = line[0].trim();
let tag_replace_str = line[1].trim();
ret.insert(tag_full_str.to_owned(), tag_replace_str.to_owned());
});
@@ -134,8 +132,14 @@ pub fn insert_message(detect_info: DetectInfo, event_time: DateTime<Utc>) {
}
/// メッセージを設定
pub fn insert(event_record: &Value, output: String, mut detect_info: DetectInfo) {
let parsed_detail = parse_message(event_record, output)
pub fn insert(
event_record: &Value,
output: String,
mut detect_info: DetectInfo,
time: DateTime<Utc>,
profile_converter: &mut HashMap<String, String>,
) {
let parsed_detail = parse_message(event_record, &output)
.chars()
.filter(|&c| !c.is_control())
.collect::<String>();
@@ -145,100 +149,41 @@ pub fn insert(event_record: &Value, output: String, mut detect_info: DetectInfo)
} else {
parsed_detail
};
let default_time = Utc.ymd(1970, 1, 1).and_hms(0, 0, 0);
let time = get_event_time(event_record).unwrap_or(default_time);
let reserverd_by_profile = _create_config_reserved_info(&detect_info, time);
for (k, v) in detect_info.ext_field.clone() {
let converted_reserve_info = convert_profile_reserved_info(v, &reserverd_by_profile);
detect_info
.ext_field
.insert(k, parse_message(event_record, converted_reserve_info));
}
insert_message(detect_info, time)
}
/// profileで用いられる予約語の情報を保持したHashMapを返す関数
fn _create_config_reserved_info(
detect_info: &DetectInfo,
time: DateTime<Utc>,
) -> HashMap<String, String> {
let mut config_reserved_info: HashMap<String, String> = HashMap::new();
for k in detect_info.ext_field.values() {
let tmp = k.as_str();
match tmp {
"%Timestamp%" => {
config_reserved_info.insert("%Timestamp%".to_string(), format_time(&time, false));
}
"%Computer%" => {
config_reserved_info.insert(
"%Computer%".to_string(),
detect_info.computername.to_owned(),
);
}
"%Details%" => {
config_reserved_info.insert("%Details%".to_string(), detect_info.detail.to_owned());
}
"%Channel%" => {
config_reserved_info
.insert("%Channel%".to_string(), detect_info.channel.to_owned());
}
"%Level%" => {
config_reserved_info.insert("%Level%".to_string(), detect_info.level.to_owned());
}
"%EventID%" => {
config_reserved_info
.insert("%EventID%".to_string(), detect_info.eventid.to_owned());
}
"%MitreAttack%" => {
config_reserved_info
.insert("%MitreAttack%".to_string(), detect_info.tag_info.to_owned());
}
"%RecordID%" => {
config_reserved_info.insert(
"%RecordID%".to_string(),
detect_info.record_id.as_ref().unwrap_or(&"-".to_string()).to_owned(),
);
}
"%RuleTitle%" => {
config_reserved_info
.insert("%RuleTitle%".to_string(), detect_info.alert.to_owned());
}
"%RecordInformation%" => {
config_reserved_info.insert(
"%RecordInformation%".to_string(),
detect_info.record_information.as_ref().unwrap_or(&"-".to_string()).to_owned(),
);
}
"%RuleFile%" => {
config_reserved_info
.insert("%RuleFile%".to_string(), detect_info.rulepath.to_owned());
}
"%EvtxFile%" => {
config_reserved_info
.insert("%EvtxFile%".to_string(), detect_info.filepath.to_owned());
}
_ => {}
let mut exist_detail = false;
PROFILES.as_ref().unwrap().iter().for_each(|(_k, v)| {
if v.contains("%Details%") {
exist_detail = true;
}
});
if exist_detail {
profile_converter.insert("%Details%".to_string(), detect_info.detail.to_owned());
}
config_reserved_info
let mut converted_detect_info = detect_info.clone();
for (k, v) in &detect_info.ext_field {
let converted_reserve_info = convert_profile_reserved_info(v, profile_converter);
converted_detect_info.ext_field.insert(
k.to_owned(),
parse_message(event_record, &converted_reserve_info),
);
}
insert_message(converted_detect_info, time)
}
/// profileで用いられる予約語の情報を変換する関数
fn convert_profile_reserved_info(
output: String,
output: &String,
config_reserved_info: &HashMap<String, String>,
) -> String {
let mut ret = output;
config_reserved_info.into_iter().for_each(|(k, v)| {
let mut ret = output.to_owned();
config_reserved_info.iter().for_each(|(k, v)| {
ret = ret.replace(k, v);
});
ret
}
/// メッセージ内の%で囲まれた箇所をエイリアスとしてをレコード情報を参照して置き換える関数
fn parse_message(event_record: &Value, output: String) -> String {
let mut return_message: String = output;
fn parse_message(event_record: &Value, output: &String) -> String {
let mut return_message = output.to_owned();
let mut hash_map: HashMap<String, String> = HashMap::new();
for caps in ALIASREGEX.captures_iter(&return_message) {
let full_target_str = &caps[0];
@@ -253,7 +198,7 @@ fn parse_message(event_record: &Value, output: String) -> String {
{
_array_str.to_string()
} else {
"Event.EventData.".to_owned() + &target_str
format!("Event.EventData.{}", target_str)
};
let split: Vec<&str> = array_str.split('.').collect();
@@ -291,7 +236,6 @@ fn parse_message(event_record: &Value, output: String) -> String {
for (k, v) in &hash_map {
return_message = return_message.replace(k, v);
}
return_message
}
@@ -472,7 +416,7 @@ mod tests {
assert_eq!(
parse_message(
&event_record,
"commandline:%CommandLine% computername:%ComputerName%".to_owned()
&"commandline:%CommandLine% computername:%ComputerName%".to_owned()
),
expected,
);
@@ -493,7 +437,7 @@ mod tests {
let event_record: Value = serde_json::from_str(json_str).unwrap();
let expected = "alias:no_alias";
assert_eq!(
parse_message(&event_record, "alias:%NoAlias%".to_owned()),
parse_message(&event_record, &"alias:%NoAlias%".to_owned()),
expected,
);
}
@@ -519,7 +463,7 @@ mod tests {
let event_record: Value = serde_json::from_str(json_str).unwrap();
let expected = "NoExistAlias:n/a";
assert_eq!(
parse_message(&event_record, "NoExistAlias:%NoAliasNoHit%".to_owned()),
parse_message(&event_record, &"NoExistAlias:%NoAliasNoHit%".to_owned()),
expected,
);
}
@@ -546,7 +490,7 @@ mod tests {
assert_eq!(
parse_message(
&event_record,
"commandline:%CommandLine% computername:%ComputerName%".to_owned()
&"commandline:%CommandLine% computername:%ComputerName%".to_owned()
),
expected,
);
@@ -579,7 +523,7 @@ mod tests {
assert_eq!(
parse_message(
&event_record,
"commandline:%CommandLine% data:%Data%".to_owned()
&"commandline:%CommandLine% data:%Data%".to_owned()
),
expected,
);
@@ -612,7 +556,7 @@ mod tests {
assert_eq!(
parse_message(
&event_record,
"commandline:%CommandLine% data:%Data[2]%".to_owned()
&"commandline:%CommandLine% data:%Data[2]%".to_owned()
),
expected,
);
@@ -645,7 +589,7 @@ mod tests {
assert_eq!(
parse_message(
&event_record,
"commandline:%CommandLine% data:%Data[0]%".to_owned()
&"commandline:%CommandLine% data:%Data[0]%".to_owned()
),
expected,
);

View File

@@ -4,6 +4,7 @@ use crate::detections::utils::check_setting_path;
use crate::yaml;
use lazy_static::lazy_static;
use linked_hash_map::LinkedHashMap;
use regex::RegexSet;
use std::collections::HashSet;
use std::fs::OpenOptions;
use std::io::{BufWriter, Write};
@@ -29,6 +30,20 @@ lazy_static! {
.values()
.cloned()
);
pub static ref PRELOAD_PROFILE: Vec<&'static str> = vec![
"%Timestamp%",
"%Computer%",
"%Channel%",
"%Level%",
"%EventID%",
"%MitreAttack%",
"%RecordID%",
"%RuleTitle%",
"%RecordInformation%",
"%RuleFile%",
"%EvtxFile%"
];
pub static ref PRELOAD_PROFILE_REGEX: RegexSet = RegexSet::new(&*PRELOAD_PROFILE).unwrap();
}
// 指定されたパスのprofileを読み込む処理
@@ -52,20 +67,14 @@ pub fn load_profile(
default_profile_path: &str,
profile_path: &str,
) -> Option<LinkedHashMap<String, String>> {
if configs::CONFIG
.read()
.unwrap()
.args
.set_default_profile
.is_some()
{
let conf = &configs::CONFIG.read().unwrap().args;
if conf.set_default_profile.is_some() {
if let Err(e) = set_default_profile(default_profile_path, profile_path) {
AlertMessage::alert(&e).ok();
} else {
println!("Successed set default profile");
};
}
let conf = &configs::CONFIG.read().unwrap().args;
let profile_all: Vec<Yaml> = if conf.profile.is_none() {
match read_profile_data(default_profile_path) {
Ok(data) => data,
@@ -91,9 +100,9 @@ pub fn load_profile(
let profile_data = &profile_all[0];
let mut ret: LinkedHashMap<String, String> = 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()
let target_data = &profile_data[profile_name.as_str()];
if !target_data.is_badvalue() {
target_data
.as_hash()
.unwrap()
.into_iter()
@@ -109,8 +118,7 @@ pub fn load_profile(
None
}
} else {
profile_all[0]
.clone()
profile_data
.as_hash()
.unwrap()
.into_iter()