From ecc1db5226f6cd5fee2bb52da5fc0061b43ed0e7 Mon Sep 17 00:00:00 2001
From: Cyrille Bagard <nocbos@gmail.com>
Date: Mon, 16 Oct 2023 01:21:31 +0200
Subject: Allow many arguments for the count() function.

---
 src/analysis/scan/items/count.c  | 17 ++++++++++++-----
 tests/analysis/scan/functions.py | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/src/analysis/scan/items/count.c b/src/analysis/scan/items/count.c
index 339c490..1d01867 100644
--- a/src/analysis/scan/items/count.c
+++ b/src/analysis/scan/items/count.c
@@ -225,17 +225,24 @@ static char *g_scan_count_function_get_name(const GScanCountFunction *item)
 static bool g_scan_count_function_run_call(GScanCountFunction *item, GScanExpression **args, size_t count, GScanContext *ctx, GScanScope *scope, GObject **out)
 {
     bool result;                            /* Bilan à retourner           */
+    size_t sum;                             /* Somme des décomptes         */
+    size_t i;                               /* Boucle de parcours          */
     size_t value;                           /* Nouveau décompte            */
 
-    if (count != 1)
-        result = false;
+    result = (count > 0);
 
-    else
+    if (result)
     {
-        result = g_scan_expression_count_items(args[0], ctx, &value);
+        sum = 0;
+
+        for (i = 0; i < count && result; i++)
+        {
+            result = g_scan_expression_count_items(args[i], ctx, &value);
+            sum += value;
+        }
 
         if (result)
-            *out = G_OBJECT(g_scan_literal_expression_new(LVT_UNSIGNED_INTEGER, (unsigned long long []){ value }));
+            *out = G_OBJECT(g_scan_literal_expression_new(LVT_UNSIGNED_INTEGER, (unsigned long long []){ sum }));
 
     }
 
diff --git a/tests/analysis/scan/functions.py b/tests/analysis/scan/functions.py
index 983b8da..6aca957 100644
--- a/tests/analysis/scan/functions.py
+++ b/tests/analysis/scan/functions.py
@@ -9,6 +9,41 @@ class TestRostFunctions(RostTestClass):
     # Core
     # ====
 
+    def testSetCounter(self):
+        """Count quantities and set sizes."""
+
+        rule = '''
+rule test {
+
+   condition:
+      count("ABC") == 3
+      and count("AB", "C") == count("ABC")
+
+}
+'''
+
+        self.check_rule_success(rule)
+
+
+        cnt = MemoryContent(b'\x01\x02\x02\x03\x03\x03')
+
+        rule = '''
+rule test {
+
+    bytes:
+        $int_01 = "\x01"
+        $int_02 = "\x02"
+        $int_3 = "\x03"
+
+    condition:
+        count($int_0*, $int_3) == #int_*
+
+}
+'''
+
+        self.check_rule_success(rule, cnt)
+
+
     def testDatasize(self):
         """Handle the size of the provided data."""
 
-- 
cgit v0.11.2-87-g4458