refactoring

This commit is contained in:
ichiichi11
2020-11-30 21:17:30 +09:00
parent cd8948de4c
commit 0e3d2ebaf4
4 changed files with 154 additions and 94 deletions

View File

@@ -1,6 +1,6 @@
use crate::detections::configs; use crate::detections::configs;
use crate::detections::print; use crate::detections::print;
use chrono::{DateTime, TimeZone, Utc}; use chrono::{DateTime, Utc};
use serde::Serialize; use serde::Serialize;
use std::error::Error; use std::error::Error;
use std::process; use std::process;
@@ -37,27 +37,37 @@ fn emit_csv(path: &str) -> Result<(), Box<dyn Error>> {
Ok(()) Ok(())
} }
#[cfg(test)]
mod tests {
use crate::afterfact::emit_csv;
use crate::detections::print;
use serde_json::Value; use serde_json::Value;
use std::fs::{read_to_string, remove_file}; use std::fs::{read_to_string, remove_file};
use std::io::Read;
#[test] #[test]
fn test_emit_csv() { fn test_emit_csv() {
{ {
let mut messages = print::MESSAGES.lock().unwrap(); let mut messages = print::MESSAGES.lock().unwrap();
let poke = Utc.ymd(1996, 2, 27).and_hms(1, 5, 1); let json_str = r##"
let json_str = r#"
{ {
"Event": { "Event": {
"EventData": { "EventData": {
"CommandLine": "hoge" "CommandLine": "hoge"
},
"System": {
"TimeCreated": {
"#attributes":{
"SystemTime": "1996-02-27T01:05:01Z"
} }
} }
} }
"#; }
}
"##;
let event_record: Value = serde_json::from_str(json_str).unwrap(); let event_record: Value = serde_json::from_str(json_str).unwrap();
messages.insert(Some(poke), &event_record, Some("pokepoke".to_string())); messages.insert(&event_record, "pokepoke".to_string());
} }
let expect = "Time,Message let expect = "Time,Message
@@ -73,3 +83,4 @@ fn test_emit_csv() {
assert!(remove_file("./test_emit_csv.csv").is_ok()); assert!(remove_file("./test_emit_csv.csv").is_ok());
} }
}

View File

@@ -88,17 +88,9 @@ impl Detection {
return; return;
} }
let event_time = Detection::get_event_time(event_record);
let utc_event_time = event_time
.and_then(|datetime| {
let utc = Utc.from_local_datetime(&datetime.naive_utc()).unwrap();
return Option::Some(utc);
})
.or(Option::None);
message.insert( message.insert(
utc_event_time,
event_record, event_record,
Some(rule.yaml["output"].as_str().unwrap().to_string()), rule.yaml["output"].as_str().unwrap_or("").to_string(),
) )
}); });
}); });
@@ -106,15 +98,4 @@ impl Detection {
// output message // output message
message.print(); message.print();
} }
fn get_event_time(event_record: &Value) -> Option<DateTime<FixedOffset>> {
let system_time =
&event_record["Event"]["System"]["TimeCreated"]["#attributes"]["SystemTime"];
let system_time_str = system_time.as_str().unwrap_or("");
if system_time_str.is_empty() {
return Option::None;
}
return DateTime::parse_from_rfc3339(system_time_str).ok();
}
} }

View File

@@ -25,39 +25,28 @@ impl Message {
} }
/// メッセージを設定 /// メッセージを設定
pub fn insert( pub fn insert(&mut self, event_record: &Value, output: String) {
&mut self, if output.is_empty() {
mut time: Option<DateTime<Utc>>,
event_record: &Value,
output: Option<String>,
) {
if Option::None == output {
return; return;
} }
let message = &self.parse_message(event_record, output); let message = &self.parse_message(event_record, output);
let default_time = Utc.ymd(1970, 1, 1).and_hms(0, 0, 0);
let time = Message::get_event_time(event_record).unwrap_or(default_time);
if Option::None == time { match self.map.get_mut(&time) {
time = Option::Some(Utc.ymd(1970, 1, 1).and_hms(0, 0, 0));
}
match self.map.get_mut(&time.unwrap()) {
Some(v) => { Some(v) => {
v.push(message.to_string()); v.push(message.to_string());
} }
None => { None => {
let m = vec![message.to_string(); 1]; let m = vec![message.to_string(); 1];
self.map.insert(time.unwrap(), m); self.map.insert(time, m);
} }
} }
} }
fn parse_message(&mut self, event_record: &Value, output: Option<String>) -> String { fn parse_message(&mut self, event_record: &Value, output: String) -> String {
if Option::None == output { let mut return_message: String = output;
return "".to_string();
}
let mut return_message: String = output.unwrap();
let mut hash_map: HashMap<String, String> = HashMap::new(); let mut hash_map: HashMap<String, String> = HashMap::new();
let re = Regex::new(r"%[a-zA-Z0-9-_]+%").unwrap(); let re = Regex::new(r"%[a-zA-Z0-9-_]+%").unwrap();
for caps in re.captures_iter(&return_message) { for caps in re.captures_iter(&return_message) {
@@ -119,15 +108,96 @@ impl Message {
pub fn iter(&self) -> &BTreeMap<DateTime<Utc>, Vec<String>> { pub fn iter(&self) -> &BTreeMap<DateTime<Utc>, Vec<String>> {
&self.map &self.map
} }
fn get_event_time(event_record: &Value) -> Option<DateTime<Utc>> {
let system_time =
&event_record["Event"]["System"]["TimeCreated"]["#attributes"]["SystemTime"];
let system_time_str = system_time.as_str().unwrap_or("");
if system_time_str.is_empty() {
return Option::None;
} }
let rfc3339_time = DateTime::parse_from_rfc3339(system_time_str);
if rfc3339_time.is_err() {
return Option::None;
}
let datetime = Utc
.from_local_datetime(&rfc3339_time.unwrap().naive_utc())
.single();
if datetime.is_none() {
return Option::None;
} else {
return Option::Some(datetime.unwrap());
}
}
}
#[cfg(test)]
mod tests {
use crate::detections::print::Message;
use serde_json::Value;
#[test] #[test]
fn test_create_and_append_message() { fn test_create_and_append_message() {
let mut message = Message::new(); let mut message = Message::new();
let poke = Utc.ymd(1996, 2, 27).and_hms(1, 5, 1); let json_str_1 = r##"
let taka = Utc.ymd(2000, 1, 21).and_hms(9, 6, 1); {
"Event": {
"EventData": {
"CommandLine": "hoge"
},
"System": {
"TimeCreated": {
"#attributes":{
"SystemTime": "1996-02-27T01:05:01Z"
}
}
}
}
}
"##;
let event_record_1: Value = serde_json::from_str(json_str_1).unwrap();
message.insert(&event_record_1, "CommandLine1: %CommandLine%".to_string());
let json_str = r#" let json_str_2 = r##"
{
"Event": {
"EventData": {
"CommandLine": "hoge"
},
"System": {
"TimeCreated": {
"#attributes":{
"SystemTime": "1996-02-27T01:05:01Z"
}
}
}
}
}
"##;
let event_record_2: Value = serde_json::from_str(json_str_2).unwrap();
message.insert(&event_record_2, "CommandLine2: %CommandLine%".to_string());
let json_str_3 = r##"
{
"Event": {
"EventData": {
"CommandLine": "hoge"
},
"System": {
"TimeCreated": {
"#attributes":{
"SystemTime": "2000-01-21T09:06:01Z"
}
}
}
}
}
"##;
let event_record_3: Value = serde_json::from_str(json_str_3).unwrap();
message.insert(&event_record_3, "CommandLine3: %CommandLine%".to_string());
let json_str_4 = r##"
{ {
"Event": { "Event": {
"EventData": { "EventData": {
@@ -135,31 +205,13 @@ fn test_create_and_append_message() {
} }
} }
} }
"#; "##;
let event_record: Value = serde_json::from_str(json_str).unwrap(); let event_record_4: Value = serde_json::from_str(json_str_4).unwrap();
message.insert(&event_record_4, "CommandLine4: %CommandLine%".to_string());
message.insert(
Some(poke),
&event_record,
Some("CommandLine1: %CommandLine%".to_string()),
);
message.insert(
Some(poke),
&event_record,
Some("CommandLine2: %CommandLine%".to_string()),
);
message.insert(
Some(taka),
&event_record,
Some("CommandLine3: %CommandLine%".to_string()),
);
message.insert(
Option::None,
&event_record,
Some("CommandLine4: %CommandLine%".to_string()),
);
let display = format!("{}", format_args!("{:?}", message)); let display = format!("{}", format_args!("{:?}", message));
println!("display::::{}", display);
let expect = "Message { map: {1970-01-01T00:00:00Z: [\"CommandLine4: hoge\"], 1996-02-27T01:05:01Z: [\"CommandLine1: hoge\", \"CommandLine2: hoge\"], 2000-01-21T09:06:01Z: [\"CommandLine3: hoge\"]} }"; let expect = "Message { map: {1970-01-01T00:00:00Z: [\"CommandLine4: hoge\"], 1996-02-27T01:05:01Z: [\"CommandLine1: hoge\", \"CommandLine2: hoge\"], 2000-01-21T09:06:01Z: [\"CommandLine3: hoge\"]} }";
assert_eq!(display, expect); assert_eq!(display, expect);
} }
}

View File

@@ -82,11 +82,27 @@ pub struct RuleNode {
impl RuleNode { impl RuleNode {
pub fn init(&mut self) -> Result<(), Vec<String>> { pub fn init(&mut self) -> Result<(), Vec<String>> {
if self.detection.is_none() { let mut errmsgs: Vec<String> = vec![];
return Result::Ok(());
// field check
if self.yaml["output"].as_str().is_none() {
errmsgs.push("Cannot find required key. key:output".to_string());
} }
return self.detection.as_mut().unwrap().init(); // detection node initialization
self.detection.as_mut().and_then(|detection| {
let detection_result = detection.init();
if detection_result.is_err() {
errmsgs.extend(detection_result.unwrap_err());
}
return Option::Some(detection);
});
if errmsgs.is_empty() {
return Result::Ok(());
} else {
return Result::Err(errmsgs);
}
} }
pub fn select(&mut self, event_record: &Value) -> bool { pub fn select(&mut self, event_record: &Value) -> bool {