fix: onetime walk with enricher (#154)

* fix: onetime_walk with enricher

* fix: delete unnecessary checking for instance type

* fix: fix unit tests

* fix: add logs

* fix: add checking for types in sort method

* fix: delete unnecessary _return_mib_enricher_for_walk method

* fix: update unit tests and change metric structure when no mib server specified

* fix: update logging messages

Author: omrozowicz-splunk

splunk_connect_for_snmp_poller/manager/task_utilities.py

@@ -85,7 +85,7 @@ async def get_translated_string(mib_server_url, var_binds, return_multimetric=Fa
     @return is_metric: boolean, metric data flag
     """
     logger.debug(f"Getting translation for the following var_binds: {var_binds}")
-    is_metric, result = await result_without_translation(var_binds)
+    is_metric, result = await result_without_translation(var_binds, return_multimetric)
 
     # Override the var_binds string with translated var_binds string
     try:
@@ -110,15 +110,14 @@ async def get_translated_string(mib_server_url, var_binds, return_multimetric=Fa
     return result, is_metric
 
 
-async def result_without_translation(var_binds):
+async def result_without_translation(var_binds, return_multimetric):
     # Get Original var_binds as backup in case the mib-server is unreachable
     for name, val in var_binds:
         # Original oid
         # TODO Discuss: should we return the original oid
         # if the mib server is unreachable
         # should we format it align with the format of the translated one
         # result = "{} = {}".format(name.prettyPrint(), val.prettyPrint())
-
         # check if this is metric data
         is_metric = is_metric_data(val.prettyPrint())
         if is_metric:
@@ -131,24 +130,29 @@ async def result_without_translation(var_binds):
             }
             result = json.dumps(result)
         else:
-            metric_content_dict = {
-                name.prettyPrint(): val.prettyPrint(),
-                InterfaceMib.METRIC_NAME_KEY: name.prettyPrint(),
-            }
+            if return_multimetric:
+                metric_content_dict = {
+                    name.prettyPrint(): val.prettyPrint(),
+                    InterfaceMib.METRIC_NAME_KEY: name.prettyPrint(),
+                }
 
-            metric_content = json.dumps(metric_content_dict)
+                metric_content = json.dumps(metric_content_dict)
 
-            non_metric_content = '{oid}="{value}"'.format(
-                oid=name.prettyPrint(), value=val.prettyPrint()
-            )
+                non_metric_content = '{oid}="{value}"'.format(
+                    oid=name.prettyPrint(), value=val.prettyPrint()
+                )
 
-            result_dict = {
-                "metric": metric_content,
-                "non_metric": non_metric_content,
-                "metric_name": name.prettyPrint(),
-            }
+                result_dict = {
+                    "metric": metric_content,
+                    "non_metric": non_metric_content,
+                    "metric_name": name.prettyPrint(),
+                }
 
-            result = json.dumps(result_dict)
+                result = json.dumps(result_dict)
+            else:
+                result = '{oid}="{value}"'.format(
+                    oid=name.prettyPrint(), value=val.prettyPrint()
+                )
         logger.debug("Our result is_metric - %s and string - %s", is_metric, result)
     return is_metric, result
 
@@ -552,48 +556,42 @@ async def walk_handler_with_enricher(
             errorIndex,
             host,
             index,
-            ir,
-            additional_metric_fields,
             one_time_flag,
             is_metric,
+            ir,
+            additional_metric_fields,
             var_binds,
         ):
             break
         else:
             result, is_metric = await get_translated_string(
                 mib_server_url, var_binds, True
             )
-            _sort_walk_data(
+            new_result = _sort_walk_data(
                 is_metric,
                 merged_result_metric,
                 merged_result_non_metric,
                 merged_result,
                 result,
             )
+            post_data_to_splunk_hec(
+                hec_sender,
+                host,
+                new_result,
+                is_metric,
+                index,
+                ir,
+                additional_metric_fields,
+                one_time_flag,
+            )
 
+    logger.info(f"Walk finished for {host} profile={ir.profile}")
     processed_result = extract_network_interface_data_from_walk(enricher, merged_result)
     additional_enricher_varbinds = (
         extract_network_interface_data_from_additional_config(enricher)
     )
-    mib_enricher = _return_mib_enricher_for_walk(
-        mongo_connection,
-        f"{host}:{port}",
-        processed_result,
-        additional_enricher_varbinds,
-    )
-    post_walk_data_to_splunk_arguments = [
-        host,
-        index,
-        one_time_flag,
-        ir,
-        additional_metric_fields,
-        mib_enricher,
-    ]
-    _post_walk_data_to_splunk(
-        hec_sender, merged_result_metric, True, *post_walk_data_to_splunk_arguments
-    )
-    _post_walk_data_to_splunk(
-        hec_sender, merged_result_non_metric, False, *post_walk_data_to_splunk_arguments
+    mongo_connection.update_mib_static_data_for(
+        f"{host}:{port}", processed_result, additional_enricher_varbinds
     )
 
 
@@ -620,56 +618,17 @@ def _sort_walk_data(
     """
     if is_metric:
         merged_result_metric.append(varbind)
+        result_to_send_to_hec = varbind
         merged_result.append(eval(varbind))
     else:
         merged_result_non_metric.append(varbind)
-        result = eval(varbind)
-        merged_result.append(eval(result["metric"]))
-
-
-def _return_mib_enricher_for_walk(
-    mongo_connection, hostname, existing_data, additional_data
-):
-    """
-    This function works only when an enricher is specified in the config and walk is being ran.
-    If any data was derived from walk result, then the function updates MongoDB with the result.
-    If no data was derived from the walk, then it's being retrieved from the MongoDB.
-    """
-    if existing_data or additional_data:
-        processed_result = mongo_connection.update_mib_static_data_for(
-            hostname, existing_data, additional_data
-        )
-        return MibEnricher(processed_result)
-    else:
-        processed_data = mongo_connection.static_data_for(hostname)
-        if not processed_data:
-            return None
-        return MibEnricher(processed_data)
-
-
-def _post_walk_data_to_splunk(
-    hec_sender,
-    result_list,
-    is_metric,
-    host,
-    index,
-    one_time_flag,
-    ir,
-    additional_metric_fields,
-    mib_enricher,
-):
-    for result in result_list:
-        post_data_to_splunk_hec(
-            hec_sender,
-            host,
-            result,
-            is_metric,
-            index,
-            ir,
-            additional_metric_fields,
-            one_time_flag,
-            mib_enricher,
-        )
+        result_dict = eval(varbind)
+        metric_part = result_dict["metric"]
+        if isinstance(metric_part, str):
+            metric_part = eval(metric_part)
+        merged_result.append(metric_part)
+        result_to_send_to_hec = result_dict["non_metric"]
+    return result_to_send_to_hec
 
 
 def parse_port(host):

splunk_connect_for_snmp_poller/mongo.py

@@ -187,6 +187,7 @@ def update_mib_static_data_for(self, host, existing_data, additional_data):
         static_data_dictionary = self.create_mib_static_data_mongo_structure(
             existing_data, additional_data
         )
+        logger.info(f"Updating static data for {host} with {static_data_dictionary}")
         self._walked_hosts.find_one_and_update(
             {"_id": host},
             {"$set": static_data_dictionary},

tests/test_task_utilities.py

@@ -81,7 +81,7 @@ def test__sort_walk_data_metric(self):
         }
         is_metric = True
         merged_result_metric, merged_result_non_metric, merged_result = [], [], []
-        _sort_walk_data(
+        result = _sort_walk_data(
             is_metric,
             merged_result_metric,
             merged_result_non_metric,
@@ -91,17 +91,19 @@ def test__sort_walk_data_metric(self):
         self.assertEqual(merged_result_metric, [varbind])
         self.assertEqual(merged_result, [varbind_dict])
         self.assertEqual(merged_result_non_metric, [])
+        self.assertEqual(result, varbind)
 
     def test__sort_walk_data_non_metric(self):
         varbind = '{"metric":"{\\"metric_name\\": \\"sc4snmp.IF-MIB.ifDescr_1\\", \\"_value\\": \\"lo\\", \\"metric_type\\": \\"OctetString\\"}","metric_name":"sc4snmp.IF-MIB.ifDescr_1","non_metric":"oid-type1=\\"ObjectIdentity\\" value1-type=\\"OctetString\\" 1.3.6.1.2.1.2.2.1.2.1=\\"lo\\" value1=\\"lo\\" IF-MIB::ifDescr.1=\\"lo\\" "}'  # noqa: E501
+        expected_result = """oid-type1="ObjectIdentity" value1-type="OctetString" 1.3.6.1.2.1.2.2.1.2.1="lo" value1="lo" IF-MIB::ifDescr.1="lo" """  # noqa: E501
         varbind_metric_dict = {
             "metric_name": "sc4snmp.IF-MIB.ifDescr_1",
             "_value": "lo",
             "metric_type": "OctetString",
         }
         is_metric = False
         merged_result_metric, merged_result_non_metric, merged_result = [], [], []
-        _sort_walk_data(
+        result = _sort_walk_data(
             is_metric,
             merged_result_metric,
             merged_result_non_metric,
@@ -111,6 +113,48 @@ def test__sort_walk_data_non_metric(self):
         self.assertEqual(merged_result_metric, [])
         self.assertEqual(merged_result, [varbind_metric_dict])
         self.assertEqual(merged_result_non_metric, [varbind])
+        self.assertEqual(result, expected_result)
+
+    def test__sort_walk_data_without_mib_server_metric(self):
+        varbind = '{"metric_name": "sc4snmp.SNMPv2-SMI.mib-2_2_2_1_10_2", "_value": "137003294"}'
+        varbind_metric_dict = {
+            "metric_name": "sc4snmp.SNMPv2-SMI.mib-2_2_2_1_10_2",
+            "_value": "137003294",
+        }
+        is_metric = True
+        merged_result_metric, merged_result_non_metric, merged_result = [], [], []
+        result = _sort_walk_data(
+            is_metric,
+            merged_result_metric,
+            merged_result_non_metric,
+            merged_result,
+            varbind,
+        )
+        self.assertEqual(merged_result_metric, [varbind])
+        self.assertEqual(merged_result, [varbind_metric_dict])
+        self.assertEqual(merged_result_non_metric, [])
+        self.assertEqual(result, varbind)
+
+    def test__sort_walk_data_without_mib_server_non_metric(self):
+        varbind = '{"metric": {"metric_name": "SNMPv2-SMI::mib-2.2.2.1.2.1", "_value": "lo"}, "non_metric": "SNMPv2-SMI::mib-2.2.2.1.2.1=lo"}'  # noqa: E501
+        varbind_metric_dict = {
+            "metric_name": "SNMPv2-SMI::mib-2.2.2.1.2.1",
+            "_value": "lo",
+        }
+        expected_result = "SNMPv2-SMI::mib-2.2.2.1.2.1=lo"
+        is_metric = False
+        merged_result_metric, merged_result_non_metric, merged_result = [], [], []
+        result = _sort_walk_data(
+            is_metric,
+            merged_result_metric,
+            merged_result_non_metric,
+            merged_result,
+            varbind,
+        )
+        self.assertEqual(merged_result_metric, [])
+        self.assertEqual(merged_result, [varbind_metric_dict])
+        self.assertEqual(merged_result_non_metric, [varbind])
+        self.assertEqual(result, expected_result)
 
     def test_is_oid_asterisk(self):
         oid = "1.3.6.1.2.1.2.2.1.1.*"