diff --git a/src/afterfact.rs b/src/afterfact.rs index 2289753ad..b9912f31e 100644 --- a/src/afterfact.rs +++ b/src/afterfact.rs @@ -30,11 +30,13 @@ use crate::detections::configs::{ }; use crate::detections::message::{AlertMessage, DetectInfo, COMPUTER_MITRE_ATTCK_MAP, LEVEL_FULL}; use crate::detections::utils::{ - self, format_time, get_writable_color, output_and_data_stack_for_html, write_color_buffer, + self, format_time, get_writable_color, output_and_data_stack_for_html, parse_csv, + write_color_buffer, }; use crate::options::htmlreport; use crate::options::profile::Profile; use crate::yaml::ParseYaml; +use rust_embed::Embed; lazy_static! { // ここで字句解析するときに使う正規表現の一覧を定義する。 @@ -48,6 +50,11 @@ lazy_static! { ]); } +#[derive(Embed)] +#[folder = "config"] +#[include = "level_color.txt"] +struct LevelColor; + #[derive(Debug)] pub struct Colors { pub output_color: termcolor::Color, @@ -976,16 +983,23 @@ pub fn get_duplicate_idxes(detect_infos: &mut [DetectInfo]) -> HashSet { /// level_color.txtファイルを読み込み対応する文字色のマッピングを返却する関数 pub fn create_output_color_map(no_color_flag: bool) -> HashMap { - let read_result = utils::read_csv( - utils::check_setting_path( - &CURRENT_EXE_PATH.to_path_buf(), - "config/level_color.txt", - true, - ) - .unwrap() - .to_str() - .unwrap(), + let path = utils::check_setting_path( + &CURRENT_EXE_PATH.to_path_buf(), + "config/level_color.txt", + false, ); + let read_result = if let Some(config_path) = path { + utils::read_csv(config_path.to_str().unwrap()) + } else { + let level_color = LevelColor::get("level_color.txt").unwrap(); + let embed_level_color = + parse_csv(std::str::from_utf8(level_color.data.as_ref()).unwrap_or_default()); + if embed_level_color.is_empty() { + Err("Not found level_color.txt in embed resource.".to_string()) + } else { + Ok(embed_level_color) + } + }; let mut color_map: HashMap = HashMap::new(); if no_color_flag { return color_map; @@ -2133,7 +2147,7 @@ fn output_detected_rule_authors( /// 与えられたyaml_pathからauthorの名前を抽出して配列で返却する関数 fn extract_author_name(yaml_path: &str) -> Nested { - let contents = match ParseYaml::read_file(Path::new(&yaml_path).to_path_buf()) { + let contents = match ParseYaml::read_file(&Path::new(&yaml_path).to_path_buf()) { Ok(yaml) => Some(yaml), Err(e) => { AlertMessage::alert(&e).ok(); diff --git a/src/detections/utils.rs b/src/detections/utils.rs index b57831a1c..941effc0a 100644 --- a/src/detections/utils.rs +++ b/src/detections/utils.rs @@ -34,6 +34,12 @@ use crate::options::htmlreport; use super::configs::{EventKeyAliasConfig, OutputOption, STORED_EKEY_ALIAS}; use super::detection::EvtxRecordInfo; use super::message::AlertMessage; +use rust_embed::Embed; + +#[derive(Embed)] +#[folder = "config"] +#[include = "default_profile_name.txt"] +pub struct DefaultProfileName; pub fn concat_selection_key(key_list: &Nested) -> String { return key_list @@ -166,13 +172,18 @@ pub fn read_csv(filename: &str) -> Result>, String> { return Result::Err(format!("Cannot open file. [file:{filename}]")); } let mut contents: String = String::new(); - let mut ret = Nested::>::new(); let read_res = f.unwrap().read_to_string(&mut contents); if let Err(e) = read_res { return Result::Err(e.to_string()); } - let mut rdr = csv::ReaderBuilder::new().from_reader(contents.as_bytes()); + let csv_res = parse_csv(&contents); + Result::Ok(csv_res) +} + +pub fn parse_csv(file_contents: &str) -> Nested> { + let mut ret = Nested::>::new(); + let mut rdr = csv::ReaderBuilder::new().from_reader(file_contents.as_bytes()); rdr.records().for_each(|r| { if r.is_err() { return; @@ -184,7 +195,7 @@ pub fn read_csv(filename: &str) -> Result>, String> { ret.push(v); }); - Result::Ok(ret) + ret } pub fn get_event_id_key() -> String { @@ -691,7 +702,7 @@ pub fn output_profile_name(output_option: &Option, stdout: bool) { // output profile name if let Some(profile_opt) = output_option { // default profile name check - let default_profile_name = read_to_string( + let default_profile_name = if let Ok(name) = read_to_string( check_setting_path( &CURRENT_EXE_PATH.to_path_buf(), "config/default_profile_name.txt", @@ -700,8 +711,14 @@ pub fn output_profile_name(output_option: &Option, stdout: bool) { .unwrap() .to_str() .unwrap(), - ) - .unwrap_or("n/a".into()); + ) { + name.trim().to_string() + } else { + let default_profile_name = DefaultProfileName::get("default_profile_name.txt").unwrap(); + std::str::from_utf8(default_profile_name.data.as_ref()) + .unwrap_or("n/a") + .to_string() + }; // user input profile option let profile_name = profile_opt diff --git a/src/main.rs b/src/main.rs index 1bed3e2d4..a328016a7 100644 --- a/src/main.rs +++ b/src/main.rs @@ -172,14 +172,6 @@ impl App { return; } - // 実行時のexeファイルのパスをベースに変更する必要があるためデフォルトの値であった場合はそのexeファイルと同一階層を探すようにする - if !CURRENT_EXE_PATH.join("config").exists() && !Path::new("./config").exists() { - AlertMessage::alert( - "Hayabusa could not find the config directory.\nPlease make sure that it is in the same directory as the hayabusa executable." - ) - .ok(); - return; - } // カレントディレクトリ以外からの実行の際にrules-configオプションの指定がないとエラーが発生することを防ぐための処理 if stored_static.config_path == Path::new("./rules/config") { stored_static.config_path = diff --git a/src/options/profile.rs b/src/options/profile.rs index 5f72f04a8..f0f8e862a 100644 --- a/src/options/profile.rs +++ b/src/options/profile.rs @@ -11,12 +11,18 @@ use crate::yaml; use compact_str::CompactString; use itertools::Itertools; use nested::Nested; +use rust_embed::Embed; use std::borrow::Cow; use std::fs::OpenOptions; use std::io::{BufWriter, Write}; use std::path::Path; use yaml_rust::{Yaml, YamlEmitter, YamlLoader}; +#[derive(Embed)] +#[folder = "config/"] +#[include = "default_profile.yaml"] +struct DefaultProfile; + #[derive(Eq, PartialEq, Hash, Clone, Debug)] pub enum Profile { Timestamp(Cow<'static, str>), @@ -134,15 +140,39 @@ impl From<&str> for Profile { // 指定されたパスのprofileを読み込む処理 fn read_profile_data(profile_path: &str) -> Result, String> { - if let Ok(loaded_profile) = yaml::ParseYaml::read_file(Path::new(profile_path).to_path_buf()) { + let profile_path_buf = Path::new(profile_path).to_path_buf(); + if let Ok(loaded_profile) = yaml::ParseYaml::read_file(&profile_path_buf) { match YamlLoader::load_from_str(&loaded_profile) { Ok(profile_yml) => Ok(profile_yml), Err(e) => Err(format!("Parse error: {profile_path}. {e}")), } } else { - Err(format!( - "The profile file({profile_path}) does not exist. Please check your default profile." - )) + let default_profile_name_path = DefaultProfile::get( + profile_path_buf + .file_name() + .unwrap() + .to_str() + .unwrap_or_default(), + ); + println!("dbg {:?}", profile_path_buf); + println!( + "dbg {:?}", + std::str::from_utf8(default_profile_name_path.clone().unwrap().data.as_ref()) + ); + // 通常のプロファイルファイルを読み込む場合 + if default_profile_name_path.is_some() { + match YamlLoader::load_from_str( + std::str::from_utf8(default_profile_name_path.unwrap().data.as_ref()) + .unwrap_or_default(), + ) { + Ok(profile_yml) => Ok(profile_yml), + Err(e) => Err(format!("Parse error: {profile_path}. {e}")), + } + } else { + Err(format!( + "The profile file({profile_path}) does not exist. Please check your default profile." + )) + } } } diff --git a/src/yaml.rs b/src/yaml.rs index 527c1f2d4..c2b08a902 100644 --- a/src/yaml.rs +++ b/src/yaml.rs @@ -54,7 +54,7 @@ impl ParseYaml { } } - pub fn read_file(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{} ", @@ -492,7 +492,7 @@ pub fn count_rules>( } // 個別のファイルの読み込みは即終了としない。 - let read_content = ParseYaml::read_file(path.as_ref().to_path_buf()); + let read_content = ParseYaml::read_file(&path.as_ref().to_path_buf()); if read_content.is_err() { return HashMap::default(); } @@ -548,7 +548,7 @@ pub fn count_rules>( } // 個別のファイルの読み込みは即終了としない。 - let read_content = ParseYaml::read_file(path); + let read_content = ParseYaml::read_file(&path); if read_content.is_err() { return io::Result::Ok(ret); } @@ -813,7 +813,7 @@ mod tests { #[test] fn test_read_yaml() { let path = Path::new("test_files/rules/yaml/1.yml"); - let ret = ParseYaml::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" { @@ -829,7 +829,7 @@ mod tests { #[test] fn test_failed_read_yaml() { let path = Path::new("test_files/rules/yaml/error.yml"); - let ret = ParseYaml::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()); }