chore: ADDON-80802 PSA implementation with uuid flag

Author: spanchal-crest

pytest_splunk_addon/app_test_generator.py

@@ -48,8 +48,9 @@ def __init__(self, pytest_config):
 
         store_events = self.pytest_config.getoption("store_events")
         config_path = self.pytest_config.getoption("splunk_data_generator")
+        ingest_with_uuid = self.pytest_config.getoption("ingest_with_uuid")
         sample_generator = SampleXdistGenerator(
-            self.pytest_config.getoption("splunk_app"), config_path
+            self.pytest_config.getoption("splunk_app"), ingest_with_uuid, config_path
         )
         store_sample = sample_generator.get_samples(store_events)
         self.tokenized_events = store_sample.get("tokenized_events")

pytest_splunk_addon/event_ingestors/hec_event_ingestor.py

@@ -93,6 +93,8 @@ def ingest(self, events, thread_count):
                 "event": event.event,
                 "index": event.metadata.get("index", "main"),
             }
+            if event.metadata["ingest_with_uuid"] == "true":
+                event_dict["fields"] = {"unique_identifier": event.unique_identifier}
 
             if event.metadata.get("host_type") in ("plugin", None):
                 host = event.metadata.get("host")

pytest_splunk_addon/event_ingestors/ingestor_helper.py

@@ -95,7 +95,7 @@ def ingest_events(
             thread_count (int): number of threads to use for ingestion
             store_events (bool): Boolean param for generating json files with tokenised events
         """
-        sample_generator = SampleXdistGenerator(addon_path, config_path)
+        sample_generator = SampleXdistGenerator(addon_path, ingest_meta_data["ingest_with_uuid"], config_path)
         store_sample = sample_generator.get_samples(store_events)
         tokenized_events = store_sample.get("tokenized_events")
         ingestor_dict = cls.get_consolidated_events(tokenized_events)

pytest_splunk_addon/fields_tests/test_generator.py

@@ -261,12 +261,15 @@ def generate_requirements_tests(self):
                     for field, value in requirement_fields.items()
                     if field not in exceptions
                 }
-                yield pytest.param(
-                    {
+                sample_event = {
                         "escaped_event": escaped_event,
                         "fields": requirement_fields,
                         "modinput_params": modinput_params,
-                    },
+                    }
+                if metadata["ingest_with_uuid"] == "true":
+                    sample_event["unique_identifier"] = event.unique_identifier
+                yield pytest.param(
+                    sample_event,
                     id=f"sample_name::{event.sample_name}::host::{event.metadata.get('host')}",
                 )
 

pytest_splunk_addon/fields_tests/test_templates.py

@@ -185,7 +185,15 @@ def test_requirements_fields(
             if param_value is not None:
                 basic_search += f" {param}={param_value}"
 
-        search = f"search {index_list} {basic_search} {escaped_event} | fields *"
+        if splunk_searchtime_fields_requirements.get("unique_identifier"):
+            record_property(
+            "stanza_name", splunk_searchtime_fields_requirements["unique_identifier"]
+        )
+            unique_identifier = splunk_searchtime_fields_requirements["unique_identifier"]
+            
+            search = f"search {index_list} {basic_search} unique_identifier=\"{unique_identifier}\" | fields *"
+        else:
+            search = f"search {index_list} {basic_search} {escaped_event} | fields *"
 
         self.logger.info(f"Executing the search query: {search}")
 
@@ -225,6 +233,7 @@ def test_requirements_fields(
         assert wrong_value_fields == {}, (
             f"\nNot all required fields have correct values or some fields are missing in Splunk. Wrong field values:\n{wrong_values_table}"
             f"{format_search_query_log(search)}"
+            f"Test failed for event: {escaped_event}\n"
         )
 
     @pytest.mark.splunk_searchtime_fields

pytest_splunk_addon/sample_generation/pytest_splunk_addon_data_parser.py

@@ -40,12 +40,13 @@ class PytestSplunkAddonDataParser:
 
     conf_name = " "
 
-    def __init__(self, addon_path: str, config_path: str):
+    def __init__(self, addon_path: str, config_path: str, ingest_with_uuid: str):
         self._conf_parser = conf_parser.TABConfigParser()
         self.config_path = config_path
         self._psa_data = None
         self.addon_path = addon_path
         self.match_stanzas = set()
+        self.ingest_with_uuid = ingest_with_uuid
         self._path_to_samples = self._get_path_to_samples()
 
     def _get_path_to_samples(self):
@@ -106,7 +107,7 @@ def get_sample_stanzas(self):
         results = []
         for sample_name, stanza_params in sorted(_psa_data.items()):
             sample_path = os.path.join(self._path_to_samples, sample_name)
-            results.append(SampleStanza(sample_path, stanza_params))
+            results.append(SampleStanza(sample_path, stanza_params, self.ingest_with_uuid))
         return results
 
     def _get_psa_data_stanzas(self):

pytest_splunk_addon/sample_generation/sample_event.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+import uuid
 import re
 import logging
 from ..index_tests import key_fields
@@ -67,6 +68,8 @@ def __init__(self, event_string, metadata, sample_name, requirement_test_data=No
         self.time_values = list()
         self.metadata = metadata
         self.sample_name = sample_name
+        if metadata["ingest_with_uuid"] == "true":
+            self.unique_identifier = str(uuid.uuid4())
         self.host_count = 0
         self.requirement_test_data = requirement_test_data
 

pytest_splunk_addon/sample_generation/sample_generator.py

@@ -33,18 +33,19 @@ class SampleGenerator(object):
     sample_stanzas = []
     conf_name = " "
 
-    def __init__(self, addon_path, config_path=None, process_count=4):
+    def __init__(self, addon_path, ingest_with_uuid, config_path=None, process_count=4):
         self.addon_path = addon_path
         self.process_count = process_count
         self.config_path = config_path
+        self.ingest_with_uuid = ingest_with_uuid
 
     def get_samples(self):
         """
         Generate SampleEvent object
         """
         if not SampleGenerator.sample_stanzas:
             psa_data_parser = PytestSplunkAddonDataParser(
-                self.addon_path, config_path=self.config_path
+                self.addon_path, config_path=self.config_path, ingest_with_uuid=self.ingest_with_uuid
             )
             sample_stanzas = psa_data_parser.get_sample_stanzas()
             SampleGenerator.conf_name = psa_data_parser.conf_name

pytest_splunk_addon/sample_generation/sample_stanza.py

@@ -47,12 +47,13 @@ class SampleStanza(object):
         psa_data_params (dict): Dictionary representing pytest-splunk-addon-data.conf
     """
 
-    def __init__(self, sample_path, psa_data_params):
+    def __init__(self, sample_path, psa_data_params, ingest_with_uuid):
         self.sample_path = sample_path
         self.sample_name = os.path.basename(sample_path)
         self.metadata = self._parse_meta(psa_data_params)
         self.sample_rules = list(self._parse_rules(psa_data_params, self.sample_path))
         self.input_type = self.metadata.get("input_type", "default")
+        self.metadata["ingest_with_uuid"] = ingest_with_uuid
         self.host_count = 0
 
     def get_raw_events(self):

pytest_splunk_addon/sample_generation/sample_xdist_generator.py

@@ -33,10 +33,11 @@ class SampleXdistGenerator:
         process_count (num): generate {no} process for execution
     """
 
-    def __init__(self, addon_path, config_path=None, process_count=4):
+    def __init__(self, addon_path, ingest_with_uuid, config_path=None, process_count=4):
         self.addon_path = addon_path
         self.process_count = process_count
         self.config_path = config_path
+        self.ingest_with_uuid = ingest_with_uuid
 
     def get_samples(self, store_events):
         """
@@ -67,7 +68,7 @@ def get_samples(self, store_events):
                         store_sample = pickle.load(file_obj)
                 else:
                     sample_generator = SampleGenerator(
-                        self.addon_path, self.config_path
+                        self.addon_path, self.ingest_with_uuid, self.config_path
                     )
                     tokenized_events = list(sample_generator.get_samples())
                     store_sample = {
@@ -79,7 +80,7 @@ def get_samples(self, store_events):
                     with open(file_path, "wb") as file_obj:
                         pickle.dump(store_sample, file_obj)
         else:
-            sample_generator = SampleGenerator(self.addon_path, self.config_path)
+            sample_generator = SampleGenerator(self.addon_path, self.ingest_with_uuid, self.config_path)
             tokenized_events = list(sample_generator.get_samples())
             store_sample = {
                 "conf_name": SampleGenerator.conf_name,
@@ -125,6 +126,7 @@ def store_events(self, tokenized_events):
                         "sourcetype": each_event.metadata.get("sourcetype"),
                         "timestamp_type": each_event.metadata.get("timestamp_type"),
                         "input_type": each_event.metadata.get("input_type"),
+                        "ingest_with_uuid": self.ingest_with_uuid,
                         "expected_event_count": expected_count,
                         "index": each_event.metadata.get("index", "main"),
                     },
@@ -137,14 +139,19 @@ def store_events(self, tokenized_events):
                         }
                     ],
                 }
+                if self.ingest_with_uuid == "true":
+                    tokenized_samples_dict[each_event.sample_name]["events"][0]["unique_identifier"] = each_event.unique_identifier
             else:
-                tokenized_samples_dict[each_event.sample_name]["events"].append(
-                    {
+                sample_event = {
                         "event": each_event.event,
                         "key_fields": each_event.key_fields,
                         "time_values": each_event.time_values,
                         "requirement_test_data": each_event.requirement_test_data,
                     }
+                if self.ingest_with_uuid == "true":
+                    sample_event["unique_identifier"] = each_event.unique_identifier
+                tokenized_samples_dict[each_event.sample_name]["events"].append(
+                    sample_event
                 )
 
         for sample_name, tokenized_sample in tokenized_samples_dict.items():

pytest_splunk_addon/splunk.py

@@ -48,6 +48,16 @@ def pytest_addoption(parser):
         by another process such as a ci/cd pipeline
     """
     group = parser.getgroup("splunk-addon")
+    group.addoption(
+        "--ingest-with-uuid",
+        action="store",
+        dest="ingest_with_uuid",
+        default="False",
+        help=(
+            "Type of ingesting and searching the events into Splunk "
+            "with uuid or without uuid."
+        ),
+    )
 
     group.addoption(
         "--splunk-app",
@@ -747,6 +757,7 @@ def splunk_ingest_data(request, splunk_hec_uri, sc4s, uf, splunk_events_cleanup)
             "splunk_hec_uri": splunk_hec_uri[1],
             "sc4s_host": sc4s[0],  # for sc4s
             "sc4s_port": sc4s[1][514],  # for sc4s
+            "ingest_with_uuid": request.config.getoption("ingest_with_uuid"),
         }
         thread_count = int(request.config.getoption("thread_count"))
         store_events = request.config.getoption("store_events")