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::print;
use chrono::{DateTime, TimeZone, Utc};
use chrono::{DateTime, Utc};
use serde::Serialize;
use std::error::Error;
use std::process;
@@ -37,27 +37,37 @@ fn emit_csv(path: &str) -> Result<(), Box<dyn Error>> {
Ok(())
}
#[cfg(test)]
mod tests {
use crate::afterfact::emit_csv;
use crate::detections::print;
use serde_json::Value;
use std::fs::{read_to_string, remove_file};
use std::io::Read;
#[test]
fn test_emit_csv() {
{
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": {
"EventData": {
"CommandLine": "hoge"
},
"System": {
"TimeCreated": {
"#attributes":{
"SystemTime": "1996-02-27T01:05:01Z"
}
}
}
"#;
}
}
"##;
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
@@ -73,3 +83,4 @@ fn test_emit_csv() {
assert!(remove_file("./test_emit_csv.csv").is_ok());
}
}

View File

@@ -88,17 +88,9 @@ impl Detection {
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(
utc_event_time,
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
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(
&mut self,
mut time: Option<DateTime<Utc>>,
event_record: &Value,
output: Option<String>,
) {
if Option::None == output {
pub fn insert(&mut self, event_record: &Value, output: String) {
if output.is_empty() {
return;
}
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 {
time = Option::Some(Utc.ymd(1970, 1, 1).and_hms(0, 0, 0));
}
match self.map.get_mut(&time.unwrap()) {
match self.map.get_mut(&time) {
Some(v) => {
v.push(message.to_string());
}
None => {
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 {
if Option::None == output {
return "".to_string();
}
let mut return_message: String = output.unwrap();
fn parse_message(&mut self, event_record: &Value, output: String) -> String {
let mut return_message: String = output;
let mut hash_map: HashMap<String, String> = HashMap::new();
let re = Regex::new(r"%[a-zA-Z0-9-_]+%").unwrap();
for caps in re.captures_iter(&return_message) {
@@ -119,15 +108,96 @@ impl Message {
pub fn iter(&self) -> &BTreeMap<DateTime<Utc>, Vec<String>> {
&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]
fn test_create_and_append_message() {
let mut message = Message::new();
let poke = Utc.ymd(1996, 2, 27).and_hms(1, 5, 1);
let taka = Utc.ymd(2000, 1, 21).and_hms(9, 6, 1);
let json_str_1 = r##"
{
"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": {
"EventData": {
@@ -135,31 +205,13 @@ fn test_create_and_append_message() {
}
}
}
"#;
let event_record: Value = serde_json::from_str(json_str).unwrap();
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 event_record_4: Value = serde_json::from_str(json_str_4).unwrap();
message.insert(&event_record_4, "CommandLine4: %CommandLine%".to_string());
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\"]} }";
assert_eq!(display, expect);
}
}

View File

@@ -82,11 +82,27 @@ pub struct RuleNode {
impl RuleNode {
pub fn init(&mut self) -> Result<(), Vec<String>> {
if self.detection.is_none() {
return Result::Ok(());
let mut errmsgs: Vec<String> = vec![];
// 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 {