Skip to content

Commit

Permalink
feat: support multiple key groupby in count condition
Browse files Browse the repository at this point in the history
  • Loading branch information
@fukusuket
fukusuket committed May 6, 2024
1 parent 48fa7e0 commit eb06004
Show file tree
Hide file tree
Showing 3 changed files with 92 additions and 56 deletions.
75 changes: 31 additions & 44 deletions src/detections/detection.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1001,58 +1001,47 @@ impl Detection {

///aggregation conditionのcount部分の検知出力文の文字列を返す関数
fn create_count_output(rule: &RuleNode, agg_result: &AggResult) -> CompactString {
// 条件式部分の出力
let mut ret: String = "[condition] ".to_string();
// この関数が呼び出されている段階で既にaggregation conditionは存在する前提なのでunwrap前の確認は行わない
let agg_condition = rule.get_agg_condition().unwrap();
let exist_timeframe = rule.yaml["detection"]["timeframe"].as_str().unwrap_or("") != "";
let mut ret: String = "".to_string();
// この関数が呼び出されている段階で既にaggregation conditionは存在する前提なのでagg_conditionの配列の長さは2となる
ret.push_str(
rule.yaml["detection"]["condition"]
.as_str()
.unwrap()
.split('|')
.nth(1)
.unwrap_or_default()
.trim(),
);
if exist_timeframe {
ret.push_str(" in timeframe");
}

write!(ret, " [result] count:{}", agg_result.data).ok();
let agg_condition = rule.get_agg_condition().unwrap();
write!(ret, "Count:{}", agg_result.data).ok();
if agg_condition._field_name.is_some() {
write!(
ret,
" {}:{}",
" ¦ {}:{}",
agg_condition._field_name.as_ref().unwrap(),
agg_result.field_values.join("/")
)
.ok();
}

if agg_condition._by_field_name.is_some() {
write!(
ret,
" {}:{}",
agg_condition._by_field_name.as_ref().unwrap(),
agg_result.key
)
.ok();
}

if exist_timeframe {
write!(
ret,
" timeframe:{}",
rule.yaml["detection"]["timeframe"].as_str().unwrap()
)
.ok();
let field_name = agg_condition._by_field_name.as_ref().unwrap();
if field_name.contains(',') {
write!(
ret,
" ¦ {}",
Self::zip_and_concat_strings(field_name, &agg_result.key)
)
.ok();
} else {
write!(ret, " ¦ {}:{}", field_name, agg_result.key).ok();
}
}

CompactString::from(ret)
}

fn zip_and_concat_strings(s1: &str, s2: &str) -> String {
let v1: Vec<&str> = s1.split(',').collect();
let v2: Vec<&str> = s2.split(',').collect();
v1.into_iter()
.zip(v2)
.map(|(s1, s2)| format!("{}:{}", s1, s2))
.collect::<Vec<String>>()
.join(" ¦ ")
}

pub fn print_rule_load_info(
rc: &HashMap<CompactString, u128>,
ld_rc: &HashMap<CompactString, u128>,
Expand Down Expand Up @@ -1337,7 +1326,7 @@ mod tests {
let test = rule_yaml.next().unwrap();
let mut rule_node = create_rule("testpath".to_string(), test);
rule_node.init(&create_dummy_stored_static()).ok();
let expected_output = "[condition] count() >= 1 [result] count:2";
let expected_output = "Count:2";
assert_eq!(
Detection::create_count_output(&rule_node, &agg_result),
expected_output
Expand All @@ -1364,7 +1353,7 @@ mod tests {
let test = rule_yaml.next().unwrap();
let mut rule_node = create_rule("testpath".to_string(), test);
rule_node.init(&create_dummy_stored_static()).ok();
let expected_output = "[condition] count() >= 1 [result] count:2";
let expected_output = "Count:2";
assert_eq!(
Detection::create_count_output(&rule_node, &agg_result),
expected_output
Expand Down Expand Up @@ -1392,8 +1381,7 @@ mod tests {
let test = rule_yaml.next().unwrap();
let mut rule_node = create_rule("testpath".to_string(), test);
rule_node.init(&create_dummy_stored_static()).ok();
let expected_output =
"[condition] count() >= 1 in timeframe [result] count:2 timeframe:15m";
let expected_output = "Count:2";
assert_eq!(
Detection::create_count_output(&rule_node, &agg_result),
expected_output
Expand Down Expand Up @@ -1423,7 +1411,7 @@ mod tests {
let test = rule_yaml.next().unwrap();
let mut rule_node = create_rule("testpath".to_string(), test);
rule_node.init(&create_dummy_stored_static()).ok();
let expected_output = "[condition] count(EventID) >= 1 [result] count:2 EventID:7040/9999";
let expected_output = "Count:2 ¦ EventID:7040/9999";
assert_eq!(
Detection::create_count_output(&rule_node, &agg_result),
expected_output
Expand Down Expand Up @@ -1453,7 +1441,7 @@ mod tests {
let test = rule_yaml.next().unwrap();
let mut rule_node = create_rule("testpath".to_string(), test);
rule_node.init(&create_dummy_stored_static()).ok();
let expected_output = "[condition] count(EventID) by process >= 1 [result] count:2 EventID:0000/1111 process:lsass.exe";
let expected_output = "Count:2 ¦ EventID:0000/1111 ¦ process:lsass.exe";
assert_eq!(
Detection::create_count_output(&rule_node, &agg_result),
expected_output
Expand Down Expand Up @@ -1482,8 +1470,7 @@ mod tests {
let test = rule_yaml.next().unwrap();
let mut rule_node = create_rule("testpath".to_string(), test);
rule_node.init(&create_dummy_stored_static()).ok();
let expected_output =
"[condition] count() by process >= 1 [result] count:2 process:lsass.exe";
let expected_output = "Count:2 ¦ process:lsass.exe";
assert_eq!(
Detection::create_count_output(&rule_node, &agg_result),
expected_output
Expand Down
30 changes: 29 additions & 1 deletion src/detections/rule/aggregation_parser.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ lazy_static! {
Regex::new(r"^>=").unwrap(),
Regex::new(r"^<").unwrap(),
Regex::new(r"^>").unwrap(),
Regex::new(r"^\w+").unwrap(),
Regex::new(r"^(\s*\w+\s*,)+\s*\w+|^\w+").unwrap(),
];
pub static ref RE_PIPE: Regex = Regex::new(r"\|.*").unwrap();
}
Expand Down Expand Up @@ -309,6 +309,34 @@ mod tests {
assert!(matches!(result._cmp_op, AggregationConditionToken::GT));
}

#[test]
fn test_aggegation_condition_compiler_count_by_multiple_fieilds() {
let compiler = AggegationConditionCompiler::new();
let result = compiler.compile("select1 or select2 | count() by iiibbb,aaabbb > 27");
assert!(result.is_ok());
let result = result.unwrap();
assert!(result.is_some());
let result = result.unwrap();
assert_eq!("iiibbb,aaabbb".to_string(), result._by_field_name.unwrap());
assert!(result._field_name.is_none());
assert_eq!(27, result._cmp_num);
assert!(matches!(result._cmp_op, AggregationConditionToken::GT));
}

#[test]
fn test_aggegation_condition_compiler_count_by_multiple_fieilds_with_space() {
let compiler = AggegationConditionCompiler::new();
let result = compiler.compile("select1 or select2 | count() by iiibbb, aaabbb > 27");
assert!(result.is_ok());
let result = result.unwrap();
assert!(result.is_some());
let result = result.unwrap();
assert_eq!("iiibbb, aaabbb".to_string(), result._by_field_name.unwrap());
assert!(result._field_name.is_none());
assert_eq!(27, result._cmp_num);
assert!(matches!(result._cmp_op, AggregationConditionToken::GT));
}

#[test]
fn test_aggegation_condition_compiler_count_field() {
let compiler = AggegationConditionCompiler::new();
Expand Down
43 changes: 32 additions & 11 deletions src/detections/rule/count.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ pub fn count(
quiet_errors_flag: bool,
json_input_flag: bool,
) {
let key = create_count_key(
let key: String = create_count_key(
rule,
record,
verbose_flag,
Expand Down Expand Up @@ -138,16 +138,37 @@ pub fn create_count_key(
let agg_condition = rule.get_agg_condition().unwrap();
if agg_condition._by_field_name.is_some() {
let by_field_key = agg_condition._by_field_name.as_ref().unwrap();
get_alias_value_in_record(
rule,
by_field_key,
record,
true,
verbose_flag,
quiet_errors_flag,
eventkey_alias,
)
.unwrap_or_else(|| "_".to_string())
if by_field_key.contains(',') {
let mut res = String::default();
for key in by_field_key.split(',') {
res.push_str(
&get_alias_value_in_record(
rule,
key.trim(),
record,
true,
verbose_flag,
quiet_errors_flag,
eventkey_alias,
)
.unwrap_or_else(|| "_".to_string()),
);
res.push(',');
}
res.pop();
res
} else {
get_alias_value_in_record(
rule,
by_field_key,
record,
true,
verbose_flag,
quiet_errors_flag,
eventkey_alias,
)
.unwrap_or_else(|| "_".to_string())
}
} else {
"_".to_string()
}
Expand Down

0 comments on commit eb06004

Please sign in to comment.