From c12090227e4fde0bc11121ef5ec3b87c3688d8e8 Mon Sep 17 00:00:00 2001 From: itiB Date: Wed, 7 Oct 2020 20:26:34 +0900 Subject: [PATCH] Fix: get rdr by reference --- src/detections/detection.rs | 13 +++++++++++-- src/detections/powershell.rs | 29 +++++++++++++---------------- src/detections/utils.rs | 15 +++++++++------ 3 files changed, 33 insertions(+), 24 deletions(-) diff --git a/src/detections/detection.rs b/src/detections/detection.rs index 1ad2daf5..683d62dd 100644 --- a/src/detections/detection.rs +++ b/src/detections/detection.rs @@ -2,13 +2,16 @@ extern crate quick_xml; use crate::detections::application; use crate::detections::common; +use crate::detections::powershell; use crate::detections::security; use crate::detections::system; -use crate::detections::powershell; use crate::models::event; use evtx::EvtxParser; use quick_xml::de::DeError; use std::collections::BTreeMap; +use std::fs::File; +use std::io::prelude::*; +extern crate csv; #[derive(Debug)] pub struct Detection { @@ -29,6 +32,12 @@ impl Detection { let mut application = application::Application::new(); let mut powershell = powershell::PowerShell::new(); + let mut f = File::open("whitelist.txt").expect("file not found"); + let mut contents = String::new(); + let _ = f.read_to_string(&mut contents); + + let mut rdr = csv::Reader::from_reader(contents.as_bytes()); + for record in parser.records() { match record { Ok(r) => { @@ -46,7 +55,7 @@ impl Detection { } else if channel == "Application" { &application.detection(event_id, &event.system, event_data); } else if channel == "Microsoft-Windows-PowerShell/Operational" { - &powershell.detection(event_id, &event.system, event_data); + &powershell.detection(event_id, &event.system, event_data, &mut rdr); } else { //&other.detection(); } diff --git a/src/detections/powershell.rs b/src/detections/powershell.rs index cf0e1b9e..3cc54337 100644 --- a/src/detections/powershell.rs +++ b/src/detections/powershell.rs @@ -2,8 +2,6 @@ use crate::detections::utils; use crate::models::event; use regex::Regex; use std::collections::HashMap; -use std::fs::File; -use std::io::prelude::*; extern crate csv; pub struct PowerShell {} @@ -18,15 +16,20 @@ impl PowerShell { event_id: String, _system: &event::System, event_data: HashMap, + rdr: &mut csv::Reader<&[u8]>, ) { if event_id == "4103" { - &self.execute_pipeline(&event_data); + &self.execute_pipeline(&event_data, rdr); } else if event_id == "4104" { - &self.execute_remote_command(&event_data); + &self.execute_remote_command(&event_data, rdr); } } - fn execute_pipeline(&mut self, event_data: &HashMap) { + fn execute_pipeline( + &mut self, + event_data: &HashMap, + rdr: &mut csv::Reader<&[u8]>, + ) { // パイプライン実行をしています let default = String::from(""); let commandline = event_data.get("ContextInfo").unwrap_or(&default); @@ -41,11 +44,6 @@ impl PowerShell { let temp = rm_before.replace_all(commandline, ""); let command = rm_after.replace_all(&temp, ""); - let mut f = File::open("whitelist.txt").expect("file not found"); - let mut contents = String::new(); - let _ = f.read_to_string(&mut contents); - - let rdr = csv::Reader::from_reader(contents.as_bytes()); if command != "" { utils::check_command(4103, &command, 1000, 0, &default, &default, rdr); } @@ -53,17 +51,16 @@ impl PowerShell { return; } - fn execute_remote_command(&mut self, event_data: &HashMap) { + fn execute_remote_command( + &mut self, + event_data: &HashMap, + rdr: &mut csv::Reader<&[u8]>, + ) { // リモートコマンドを実行します let default = String::from(""); let message_num = event_data.get("MessageNumber"); let commandline = event_data.get("ScriptBlockText").unwrap_or(&default); - let mut f = File::open("whitelist.txt").expect("file not found"); - let mut contents = String::new(); - let _ = f.read_to_string(&mut contents); - - let rdr = csv::Reader::from_reader(contents.as_bytes()); match message_num { Some(_) => utils::check_command(4104, &commandline, 1000, 0, &default, &default, rdr), _ => {} diff --git a/src/detections/utils.rs b/src/detections/utils.rs index 9a2ffdc3..c18378c8 100644 --- a/src/detections/utils.rs +++ b/src/detections/utils.rs @@ -16,7 +16,7 @@ pub fn check_command( servicecmd: usize, servicename: &str, creator: &str, - mut rdr: csv::Reader<&[u8]>, + rdr: &mut csv::Reader<&[u8]>, ) { let mut text = "".to_string(); let mut base64 = "".to_string(); @@ -148,7 +148,10 @@ fn check_obfu(string: &str) -> std::string::String { fn check_regex(string: &str, r#type: usize) -> std::string::String { let mut f = File::open("regexes.txt").expect("file not found"); let mut contents = String::new(); - f.read_to_string(&mut contents); + let ret = f.read_to_string(&mut contents); + if let Err(_) = ret { + return "".to_string(); + } let mut rdr = csv::Reader::from_reader(contents.as_bytes()); @@ -222,10 +225,10 @@ mod tests { let mut contents = String::new(); f.read_to_string(&mut contents); - let rdr = csv::Reader::from_reader(contents.as_bytes()); - utils::check_command(1, "dir", 100, 100, "dir", "dir", rdr); + let mut rdr = csv::Reader::from_reader(contents.as_bytes()); + utils::check_command(1, "dir", 100, 100, "dir", "dir", &mut rdr); - let rdr = csv::Reader::from_reader(contents.as_bytes()); + let mut rdr = csv::Reader::from_reader(contents.as_bytes()); //test return with whitelist. utils::check_command( 1, @@ -234,7 +237,7 @@ mod tests { 100, "dir", "dir", - rdr, + &mut rdr, ); } }