fix: fix issue with onetime walk not running (#168)

* fix: fix issue with onetime walk not running

* fix: fix unit tests

* fix: add fix for profile matching onetime IF-MIB walk

Author: omrozowicz-splunk

splunk_connect_for_snmp_poller/manager/poller.py

@@ -123,35 +123,39 @@ def __check_inventory(self):
                     continue
                 inventory_entry_keys.add(entry_key)
                 ir_host = return_database_id(ir.host)
-                if ir_host not in inventory_hosts:
-                    self.add_enricher_to_a_host(
-                        new_enricher, profiles, copy.deepcopy(ir), True
-                    )
-                inventory_hosts.add(return_database_id(ir_host))
+                inventory_hosts.add(ir_host)
                 inventory_hosts_with_snmp_data[ir_host] = copy.deepcopy(ir)
                 if ir.profile == DYNAMIC_PROFILE:
                     self.delete_all_entries_per_host(ir.host)
                     self.add_device_for_profile_matching(ir)
+                    self.check_if_new_host_was_added(entry_key, ir, new_enricher)
                 else:
                     logger.debug(
                         "[-] server_config['profiles']: %s",
                         self._server_config["profiles"],
                     )
                     if entry_key not in self._jobs_map:
+                        self.check_if_new_host_was_added(entry_key, ir, new_enricher)
                         self.process_new_job(entry_key, ir, profiles)
                     else:
                         self.update_schedule_for_changed_conf(entry_key, ir, profiles)
 
             if server_config_modified:
                 if new_enricher != self._old_enricher:
-                    self.run_enricher_check(
-                        new_enricher, profiles, inventory_hosts_with_snmp_data
+                    self.run_enricher_changed_check(
+                        new_enricher, inventory_hosts_with_snmp_data
                     )
             self.clean_job_inventory(inventory_entry_keys, inventory_hosts)
 
-    def run_enricher_check(
-        self, new_enricher, profiles, inventory_hosts_with_snmp_data
-    ):
+    def check_if_new_host_was_added(self, host_key, inventory_record, new_enricher):
+        ir_host = return_database_id(host_key)
+        if self._old_enricher != {}:
+            logger.info(f"New host: {ir_host}")
+            self.__add_enricher_to_a_host(
+                new_enricher, copy.deepcopy(inventory_record), True
+            )
+
+    def run_enricher_changed_check(self, new_enricher, inventory_hosts_with_snmp_data):
         logger.info(
             f"Previous enricher: {self._old_enricher} \n New enricher: {new_enricher}"
         )
@@ -161,19 +165,19 @@ def run_enricher_check(
             self._old_enricher = {}
             return
         for inventory_host in inventory_hosts_with_snmp_data.keys():
-            self.add_enricher_to_a_host(
-                new_enricher, profiles, inventory_hosts_with_snmp_data[inventory_host]
+            self.__add_enricher_to_a_host(
+                new_enricher, inventory_hosts_with_snmp_data[inventory_host]
             )
         self._old_enricher = new_enricher
 
-    def add_enricher_to_a_host(self, current_enricher, profiles, ir, new_host=False):
+    def __add_enricher_to_a_host(self, current_enricher, ir, new_host=False):
+        logger.info("Add enricher to a host")
         old_enricher = {} if new_host else self._old_enricher
         if current_enricher != {}:
             update_enricher_config(
                 old_enricher,
                 current_enricher,
                 self._mongo,
-                profiles,
                 ir,
                 self._server_config,
                 self.__get_splunk_indexes(),

splunk_connect_for_snmp_poller/manager/poller_utilities.py

@@ -163,7 +163,8 @@ def _extract_sys_uptime_instance(
 
 
 def _walk_info(mongo_collection, host, current_sys_up_time):
-    host_already_walked = mongo_collection.contains_host(host) != 0
+    host_already_walked = mongo_collection.first_time_walked_was_initiated(host) != 0
+    logger.info(f"host_already_walked: {host_already_walked}")
     should_do_walk = not host_already_walked
     if host_already_walked:
         previous_sys_up_time = mongo_collection.real_time_data_for(host)
@@ -275,6 +276,10 @@ def automatic_realtime_task(
                 host_already_walked,
                 sys_up_time,
             )
+            if should_do_walk:
+                mongo_collection.update_walked_host(
+                    db_host_id, {"walked_first_time": True}
+                )
     except Exception:
         logger.exception("Error during automatic_realtime_task")
 
@@ -283,15 +288,14 @@ def update_enricher_config(
     old_enricher,
     new_enricher,
     mongo,
-    profiles,
     inventory_host,
     server_config,
     splunk_indexes,
 ):
     run_ifmib_walk = is_ifmib_different(old_enricher, new_enricher)
     if run_ifmib_walk:
         _update_enricher_config_with_ifmib(
-            profiles, inventory_host, server_config, splunk_indexes
+            inventory_host, server_config, splunk_indexes
         )
     else:
         _update_enricher_config_for_additional_varbinds(
@@ -304,17 +308,17 @@ def update_enricher_config(
 
 
 def _update_enricher_config_with_ifmib(
-    profiles,
     inventory_host,
     server_config,
     splunk_indexes,
 ):
     inventory_host.profile = OidConstant.IF_MIB
-    snmp_polling.delay(
-        inventory_host.to_json(),
+    schedule.every().second.do(
+        onetime_task,
+        inventory_host,
         server_config,
         splunk_indexes,
-        profiles,
+        None,
     )
 
 

splunk_connect_for_snmp_poller/mongo.py

@@ -108,8 +108,14 @@ def is_connected(self):
     def contains_host(self, host):
         return self._walked_hosts.find({"_id": host}).count()
 
+    def first_time_walked_was_initiated(self, host):
+        return self._walked_hosts.find({"_id": host, "walked_first_time": True}).count()
+
     def add_host(self, host):
-        self._walked_hosts.insert_one({"_id": host})
+        try:
+            self._walked_hosts.insert_one({"_id": host})
+        except:  # noqa: E722
+            logger.info(f"Id {host} already exists in MongoDB")
 
     def get_all_unwalked_hosts(self):
         return list(self._unwalked_hosts.find({}))
@@ -136,6 +142,14 @@ def delete_host(self, host):
     def clear(self):
         self._walked_hosts.remove()
 
+    def update_walked_host(self, host, element):
+        self._walked_hosts.find_one_and_update(
+            {"_id": host},
+            {"$set": element},
+            return_document=ReturnDocument.AFTER,
+            upsert=True,
+        )
+
     def real_time_data_for(self, host):
         full_collection = self._walked_hosts.find_one({"_id": host})
         if (

tests/test_poller.py

@@ -22,7 +22,7 @@
 
 
 class TestPollerUtilities(TestCase):
-    def test_run_enricher_check_when_enricher_is_deleted(self):
+    def test_run_enricher_changed_check_when_enricher_is_deleted(self):
         server_config = {"mongo": ""}
         with patch(
             "splunk_connect_for_snmp_poller.mongo.WalkedHostsRepository.__init__"
@@ -42,5 +42,5 @@ def test_run_enricher_check_when_enricher_is_deleted(self):
             with patch(
                 "splunk_connect_for_snmp_poller.mongo.WalkedHostsRepository.delete_all_static_data"
             ):
-                obj.run_enricher_check({}, [], {"127.0.0.1:161": MagicMock()})
+                obj.run_enricher_changed_check({}, {"127.0.0.1:161": MagicMock()})
             self.assertEqual(obj._old_enricher, {})