refactoring

This commit is contained in:
HajimeTakai
2021-05-09 17:26:17 +09:00
parent 2f24dc775f
commit 7913fbfb95
3 changed files with 54 additions and 17 deletions

7
Cargo.lock generated
View File

@@ -632,6 +632,12 @@ dependencies = [
"winapi", "winapi",
] ]
[[package]]
name = "mopa"
version = "0.2.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "a785740271256c230f57462d3b83e52f998433a7062fc18f96d5999474a9f915"
[[package]] [[package]]
name = "ntapi" name = "ntapi"
version = "0.3.6" version = "0.3.6"
@@ -1358,6 +1364,7 @@ dependencies = [
"flate2", "flate2",
"lazy_static", "lazy_static",
"linked-hash-map", "linked-hash-map",
"mopa",
"num_cpus", "num_cpus",
"quick-xml 0.17.2", "quick-xml 0.17.2",
"regex", "regex",

View File

@@ -23,6 +23,7 @@ yaml-rust = "0.4"
linked-hash-map = "0.5.3" linked-hash-map = "0.5.3"
tokio = { version = "1", features = ["full"] } tokio = { version = "1", features = ["full"] }
num_cpus = "1.13.0" num_cpus = "1.13.0"
mopa = "0.2.2"
[target.x86_64-pc-windows-gnu] [target.x86_64-pc-windows-gnu]
linker = "x86_64-w64-mingw32-gcc" linker = "x86_64-w64-mingw32-gcc"

View File

@@ -1,5 +1,7 @@
extern crate regex; extern crate regex;
use mopa::mopafy;
use std::vec; use std::vec;
use crate::detections::utils; use crate::detections::utils;
@@ -135,16 +137,17 @@ impl RuleNode {
return selection return selection
.unwrap() .unwrap()
.get_leaf_nodes() .get_descendants()
.iter() .iter()
.filter_map(|node| return node.downcast_ref::<LeafSelectionNode>()) // mopaというライブラリを使うと簡単にダウンキャストできるらしいです。https://crates.io/crates/mopa
.filter(|node| { .filter(|node| {
// alias.txtのevent_keyに一致するかどうか // キーがEventIDのードである
let key = utils::get_event_id_key(); let key = utils::get_event_id_key();
if node.get_key() == key { if node.get_key() == key {
return true; return true;
} }
// alias.txtのaliasに一致るかどうか // EventIDのAliasに一致しているかどうか
let alias = utils::get_alias(&key); let alias = utils::get_alias(&key);
if alias.is_none() { if alias.is_none() {
return false; return false;
@@ -175,11 +178,13 @@ impl DetectionNode {
} }
// Ruleファイルの detection- selection配下のードはこのtraitを実装する。 // Ruleファイルの detection- selection配下のードはこのtraitを実装する。
trait SelectionNode { trait SelectionNode: mopa::Any {
fn select(&self, event_record: &Value) -> bool; fn select(&self, event_record: &Value) -> bool;
fn init(&mut self) -> Result<(), Vec<String>>; fn init(&mut self) -> Result<(), Vec<String>>;
fn get_leaf_nodes(&self) -> Vec<&LeafSelectionNode>; fn get_childs(&self) -> Vec<&Box<dyn SelectionNode>>;
fn get_descendants(&self) -> Vec<&Box<dyn SelectionNode>>;
} }
mopafy!(SelectionNode);
// detection - selection配下でAND条件を表すード // detection - selection配下でAND条件を表すード
struct AndSelectionNode { struct AndSelectionNode {
@@ -230,17 +235,26 @@ impl SelectionNode for AndSelectionNode {
} }
} }
fn get_leaf_nodes(&self) -> Vec<&LeafSelectionNode> { fn get_childs(&self) -> Vec<&Box<dyn SelectionNode>> {
let mut ret = vec![]; let mut ret = vec![];
self.child_nodes.iter().for_each(|child_node| {
ret.push(child_node);
});
return ret;
}
fn get_descendants(&self) -> Vec<&Box<dyn SelectionNode>> {
let mut ret = self.get_childs();
self.child_nodes self.child_nodes
.iter() .iter()
.map(|child| { .map(|child_node| {
return child.get_leaf_nodes(); return child_node.get_descendants();
}) })
.flatten() .flatten()
.for_each(|descendant| { .for_each(|descendant_node| {
ret.push(descendant); ret.push(descendant_node);
}); });
return ret; return ret;
@@ -296,17 +310,26 @@ impl SelectionNode for OrSelectionNode {
} }
} }
fn get_leaf_nodes(&self) -> Vec<&LeafSelectionNode> { fn get_childs(&self) -> Vec<&Box<dyn SelectionNode>> {
let mut ret = vec![]; let mut ret = vec![];
self.child_nodes.iter().for_each(|child_node| {
ret.push(child_node);
});
return ret;
}
fn get_descendants(&self) -> Vec<&Box<dyn SelectionNode>> {
let mut ret = self.get_childs();
self.child_nodes self.child_nodes
.iter() .iter()
.map(|child| { .map(|child_node| {
return child.get_leaf_nodes(); return child_node.get_descendants();
}) })
.flatten() .flatten()
.for_each(|descendant| { .for_each(|descendant_node| {
ret.push(descendant); ret.push(descendant_node);
}); });
return ret; return ret;
@@ -453,8 +476,12 @@ impl SelectionNode for LeafSelectionNode {
.init(&match_key_list, &self.select_value); .init(&match_key_list, &self.select_value);
} }
fn get_leaf_nodes(&self) -> Vec<&LeafSelectionNode> { fn get_childs(&self) -> Vec<&Box<dyn SelectionNode>> {
return vec![&self]; return vec![];
}
fn get_descendants(&self) -> Vec<&Box<dyn SelectionNode>> {
return vec![];
} }
} }
@@ -727,3 +754,5 @@ impl LeafMatcher for WhitelistFileMatcher {
}; };
} }
} }