Skip to content

Commit

Permalink
feat: embedded config files to binary #1370
Browse files Browse the repository at this point in the history
  • Loading branch information
@hitenkoku
hitenkoku committed Jul 15, 2024
1 parent ec8213a commit ee9e93c
Show file tree
Hide file tree
Showing 5 changed files with 89 additions and 36 deletions.
36 changes: 25 additions & 11 deletions src/afterfact.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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! {
// ここで字句解析するときに使う正規表現の一覧を定義する。
Expand All @@ -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,
Expand Down Expand Up @@ -976,16 +983,23 @@ pub fn get_duplicate_idxes(detect_infos: &mut [DetectInfo]) -> HashSet<usize> {

/// level_color.txtファイルを読み込み対応する文字色のマッピングを返却する関数
pub fn create_output_color_map(no_color_flag: bool) -> HashMap<CompactString, Colors> {
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<CompactString, Colors> = HashMap::new();
if no_color_flag {
return color_map;
Expand Down Expand Up @@ -2133,7 +2147,7 @@ fn output_detected_rule_authors(

/// 与えられたyaml_pathからauthorの名前を抽出して配列で返却する関数
fn extract_author_name(yaml_path: &str) -> Nested<String> {
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();
Expand Down
29 changes: 23 additions & 6 deletions src/detections/utils.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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>) -> String {
return key_list
Expand Down Expand Up @@ -166,13 +172,18 @@ pub fn read_csv(filename: &str) -> Result<Nested<Vec<String>>, String> {
return Result::Err(format!("Cannot open file. [file:{filename}]"));
}
let mut contents: String = String::new();
let mut ret = Nested::<Vec<String>>::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<Vec<String>> {
let mut ret = Nested::<Vec<String>>::new();
let mut rdr = csv::ReaderBuilder::new().from_reader(file_contents.as_bytes());
rdr.records().for_each(|r| {
if r.is_err() {
return;
Expand All @@ -184,7 +195,7 @@ pub fn read_csv(filename: &str) -> Result<Nested<Vec<String>>, String> {
ret.push(v);
});

Result::Ok(ret)
ret
}

pub fn get_event_id_key() -> String {
Expand Down Expand Up @@ -691,7 +702,7 @@ pub fn output_profile_name(output_option: &Option<OutputOption>, 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",
Expand All @@ -700,8 +711,14 @@ pub fn output_profile_name(output_option: &Option<OutputOption>, 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
Expand Down
8 changes: 0 additions & 8 deletions src/main.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 =
Expand Down
38 changes: 34 additions & 4 deletions src/options/profile.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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>),
Expand Down Expand Up @@ -134,15 +140,39 @@ impl From<&str> for Profile {

// 指定されたパスのprofileを読み込む処理
fn read_profile_data(profile_path: &str) -> Result<Vec<Yaml>, 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."
))
}
}
}

Expand Down
14 changes: 7 additions & 7 deletions src/yaml.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ impl ParseYaml {
}
}

pub fn read_file(path: PathBuf) -> Result<String, String> {
pub fn read_file(path: &PathBuf) -> Result<String, String> {
let mut file_content = String::new();

let mut fr = fs::File::open(path)
Expand Down Expand Up @@ -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{} ",
Expand Down Expand Up @@ -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{} ",
Expand Down Expand Up @@ -492,7 +492,7 @@ pub fn count_rules<P: AsRef<Path>>(
}

// 個別のファイルの読み込みは即終了としない。
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();
}
Expand Down Expand Up @@ -548,7 +548,7 @@ pub fn count_rules<P: AsRef<Path>>(
}

// 個別のファイルの読み込みは即終了としない。
let read_content = ParseYaml::read_file(path);
let read_content = ParseYaml::read_file(&path);
if read_content.is_err() {
return io::Result::Ok(ret);
}
Expand Down Expand Up @@ -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" {
Expand All @@ -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());
}
Expand Down

0 comments on commit ee9e93c

Please sign in to comment.