[crowdstrike] feat: migrated create traces to bulk endpoint (#2873)

Author: impolitepanda

crowdstrike/crowdstrike/openbas_crowdstrike.py

@@ -78,7 +78,23 @@ def _fetch_expectations(self, start_time):
 
     # --- MATCHING ---
 
-    def _match_expectations(self, alerts: list[Item], expectations):
+    def _match_expectations(
+        self, alerts: list[Item], expectations
+    ) -> list[dict[str, str]]:
+        """
+        Match OpenBAS expectations with CS alerts.
+
+        Args:
+            alerts: The list of alerts received from CS.
+            expectations: The list of OpenBAS expectation to match to CS alerts.
+
+        Returns:
+            * A list of expectation traces that need to be created following all the matches, if any alerts matched any
+              of the given expectations.
+              OR
+            * [], otherwise.
+        """
+        traces_to_create: list[dict[str, str]] = []
 
         self.helper.collector_logger.debug(
             "Total expectations returned: " + str(len(expectations))
@@ -131,16 +147,16 @@ def _match_expectations(self, alerts: list[Item], expectations):
                         )
                         expectations_not_filled.remove(expectation)
 
-                    # Send alert to openbas for current matched expectation. Duplicate alerts are handled by openbas itself
+                    # Save alert to openbas for current matched expectation. Duplicate alerts are handled by openbas itself
                     self.helper.collector_logger.info(
                         "Expectation matched, adding trace for expectation "
                         + expectation["inject_expectation_id"]
                         + " ("
                         + expectation["inject_expectation_type"]
                         + ")"
                     )
-                    self.helper.api.inject_expectation_trace.create(
-                        data={
+                    traces_to_create.append(
+                        {
                             "inject_expectation_trace_expectation": expectation[
                                 "inject_expectation_id"
                             ],
@@ -157,6 +173,8 @@ def _match_expectations(self, alerts: list[Item], expectations):
                         }
                     )
 
+        return traces_to_create
+
     # --- PROCESS ---
 
     def _is_expectation_filled(self, expectation) -> bool:
@@ -172,11 +190,15 @@ def _process(self):
         now = datetime.now(pytz.UTC)
         start_time = now - timedelta(minutes=self.scanning_delta)
 
-        self._match_expectations(
+        traces = self._match_expectations(
             alerts=self.strategy.get_raw_data(start_time),
             expectations=self._fetch_expectations(start_time),
         )
 
+        self.helper.api.inject_expectation_trace.bulk_create(
+            payload={"expectation_traces": traces}
+        )
+
     def start(self):
         period = self.config.get_conf("collector_period", True, 60)
         self.helper.schedule(self._process, period)

crowdstrike/test/test_openbas_crowdstrike.py

@@ -7,10 +7,9 @@
 
 
 class TestOpenBASCrowdstrike(unittest.TestCase):
-    @patch("pyobas.apis.InjectExpectationTraceManager.create")
     @patch("pyobas.apis.InjectExpectationManager.update")
     def test_when_alert_matches_update_prevention_expectation(
-        self, mock_expectation_update, mock_traces_create
+        self, mock_expectation_update
     ):
         expected_expectation_id = "expectation_id"
         expectations = [
@@ -41,7 +40,7 @@ def test_when_alert_matches_update_prevention_expectation(
 
         collector = get_default_collector(strategy)
 
-        collector._match_expectations([Item(**MOCKED_ALERT)], expectations)
+        traces = collector._match_expectations([Item(**MOCKED_ALERT)], expectations)
 
         mock_expectation_update.assert_called_once_with(
             expected_expectation_id,
@@ -52,9 +51,10 @@ def test_when_alert_matches_update_prevention_expectation(
                 "metadata": {"alertId": expected_expectation_id},
             },
         )
-
-        mock_traces_create.assert_called_once_with(
-            data={
+        self.assertEqual(1, len(traces))
+        self.assertEqual(
+            traces[0],
+            {
                 "inject_expectation_trace_expectation": expected_expectation_id,
                 "inject_expectation_trace_source_id": collector.config.get_conf(
                     "collector_id"
@@ -69,10 +69,9 @@ def test_when_alert_matches_update_prevention_expectation(
             },
         )
 
-    @patch("pyobas.apis.InjectExpectationTraceManager.create")
     @patch("pyobas.apis.InjectExpectationManager.update")
     def test_when_alert_matches_but_not_prevented_update_prevention_expectation(
-        self, mock_expectation_update, mock_traces_create
+        self, mock_expectation_update
     ):
         expected_expectation_id = "expectation_id"
         expectations = [
@@ -103,7 +102,7 @@ def test_when_alert_matches_but_not_prevented_update_prevention_expectation(
 
         collector = get_default_collector(strategy)
 
-        collector._match_expectations([Item(**MOCKED_ALERT)], expectations)
+        traces = collector._match_expectations([Item(**MOCKED_ALERT)], expectations)
 
         mock_expectation_update.assert_called_once_with(
             expected_expectation_id,
@@ -115,8 +114,10 @@ def test_when_alert_matches_but_not_prevented_update_prevention_expectation(
             },
         )
 
-        mock_traces_create.assert_called_once_with(
-            data={
+        self.assertEqual(1, len(traces))
+        self.assertEqual(
+            traces[0],
+            {
                 "inject_expectation_trace_expectation": expected_expectation_id,
                 "inject_expectation_trace_source_id": collector.config.get_conf(
                     "collector_id"
@@ -131,10 +132,9 @@ def test_when_alert_matches_but_not_prevented_update_prevention_expectation(
             },
         )
 
-    @patch("pyobas.apis.InjectExpectationTraceManager.create")
     @patch("pyobas.apis.InjectExpectationManager.update")
     def test_when_alert_matches_update_detection_expectation(
-        self, mock_expectation_update, mock_traces_create
+        self, mock_expectation_update
     ):
         expected_expectation_id = "expectation_id"
         expectations = [
@@ -165,7 +165,7 @@ def test_when_alert_matches_update_detection_expectation(
 
         collector = get_default_collector(strategy)
 
-        collector._match_expectations([Item(**MOCKED_ALERT)], expectations)
+        traces = collector._match_expectations([Item(**MOCKED_ALERT)], expectations)
 
         mock_expectation_update.assert_called_once_with(
             expected_expectation_id,
@@ -177,8 +177,10 @@ def test_when_alert_matches_update_detection_expectation(
             },
         )
 
-        mock_traces_create.assert_called_once_with(
-            data={
+        self.assertEqual(1, len(traces))
+        self.assertEqual(
+            traces[0],
+            {
                 "inject_expectation_trace_expectation": expected_expectation_id,
                 "inject_expectation_trace_source_id": collector.config.get_conf(
                     "collector_id"
@@ -193,10 +195,9 @@ def test_when_alert_matches_update_detection_expectation(
             },
         )
 
-    @patch("pyobas.apis.InjectExpectationTraceManager.create")
     @patch("pyobas.apis.InjectExpectationManager.update")
     def test_when_expectation_has_expected_hostname_signature_ignore_it(
-        self, mock_expectation_update, mock_traces_create
+        self, mock_expectation_update
     ):
         expected_expectation_id = "expectation_id"
         expectations = [
@@ -228,7 +229,7 @@ def test_when_expectation_has_expected_hostname_signature_ignore_it(
 
         collector = get_default_collector(strategy)
 
-        collector._match_expectations([Item(**MOCKED_ALERT)], expectations)
+        traces = collector._match_expectations([Item(**MOCKED_ALERT)], expectations)
 
         mock_expectation_update.assert_called_once_with(
             expected_expectation_id,
@@ -240,8 +241,10 @@ def test_when_expectation_has_expected_hostname_signature_ignore_it(
             },
         )
 
-        mock_traces_create.assert_called_once_with(
-            data={
+        self.assertEqual(1, len(traces))
+        self.assertEqual(
+            traces[0],
+            {
                 "inject_expectation_trace_expectation": expected_expectation_id,
                 "inject_expectation_trace_source_id": collector.config.get_conf(
                     "collector_id"
@@ -289,9 +292,10 @@ def test_when_alert_fails_to_match_dont_update_prevention_expectation(
 
         collector = get_default_collector(strategy)
 
-        collector._match_expectations([Item(**MOCKED_ALERT)], expectations)
+        traces = collector._match_expectations([Item(**MOCKED_ALERT)], expectations)
 
         mock_expectation_update.assert_not_called()
+        self.assertEqual(0, len(traces))
 
     @patch("pyobas.apis.InjectExpectationManager.update")
     def test_when_signatures_match_when_unknown_expectation_type_skip_updating_expectation(
@@ -326,10 +330,12 @@ def test_when_signatures_match_when_unknown_expectation_type_skip_updating_expec
 
         collector = get_default_collector(strategy)
 
-        collector._match_expectations([Item(**MOCKED_ALERT)], expectations)
+        traces = collector._match_expectations([Item(**MOCKED_ALERT)], expectations)
 
         mock_expectation_update.assert_not_called()
 
+        self.assertEqual(0, len(traces))
+
 
 if __name__ == "__main__":
     unittest.main()