Merge pull request #1584 from stratosphereips/alya/fix-latency/remove_cpu_intensive_dead_code

Remove cpu intensive dead code

Author: AlyaGomaa

modules/flowalerts/flowalerts.py

@@ -97,7 +97,7 @@ async def main(self):
                     loop = asyncio.get_event_loop()
                     task = loop.create_task(analyzer.analyze(msg))
                     # because Async Tasks swallow exceptions.
-                    task.add_done_callback(self.handle_exception)
+                    task.add_done_callback(self.handle_task_exception)
                     # to wait for these functions before flowalerts shuts down
                     self.tasks.append(task)
                     # Allow the event loop to run the scheduled task

modules/ip_info/ip_info.py

@@ -1,6 +1,5 @@
 # SPDX-FileCopyrightText: 2021 Sebastian Garcia <sebastian.garcia@agents.fel.cvut.cz>
 # SPDX-License-Identifier: GPL-2.0-only
-from asyncio import Task
 from typing import (
     Union,
     Optional,
@@ -96,7 +95,7 @@ async def open_dbs(self):
                 "https://dev.maxmind.com/geoip/geolite2-free-geolocation-data?lang=en. "
                 "Please note it must be the MaxMind DB version."
             )
-        self.reading_mac_db_task: Task = self.create_task(self.read_mac_db)
+        self.create_task(self.read_mac_db)
 
     async def read_mac_db(self):
         """
@@ -325,7 +324,6 @@ async def shutdown_gracefully(self):
             self.country_db.close()
         if hasattr(self, "mac_db"):
             self.mac_db.close()
-        await self.reading_mac_db_task
 
     # GW
     @staticmethod

modules/update_manager/update_manager.py

@@ -1653,22 +1653,14 @@ def delete_unused_cached_remote_feeds(self):
                 )
         self.loaded_ti_files -= 1
 
-    def handle_exception(self, task):
-        """
-        in asyncmodules we use Async.Task to run some of the functions
-        If an exception occurs in a coroutine that was wrapped in a Task
-        (e.g., asyncio.create_task), the exception does not crash the program
-         but remains in the task.
-        This function is used to handle the exception in the task
-        """
+    def handle_task_exception(self, task):
         try:
-            # Access task result to raise the exception if it occurred
-            task.result()
-        except asyncio.exceptions.CancelledError:
-            # like pressing ctrl+c
-            return
-        except Exception as e:
-            self.print(e, 0, 1)
+            exception = task.exception()
+        except asyncio.CancelledError:
+            return  # Task was cancelled, not an error
+        if exception:
+            self.print(f"Unhandled exception in task: {exception}")
+            self.print_traceback()
 
     async def update(self) -> bool:
         """
@@ -1725,7 +1717,7 @@ async def update(self) -> bool:
                     task = asyncio.create_task(
                         self.update_ti_file(file_to_download)
                     )
-                    task.add_done_callback(self.handle_exception)
+                    task.add_done_callback(self.handle_task_exception)
             #######################################################
             # in case of riskiq files, we don't have a link for them in ti_files, We update these files using their API
             # check if we have a username and api key and a week has passed since we last updated
@@ -1753,7 +1745,7 @@ async def update_ti_files(self):
         # create_task is used to run update() function
         # concurrently instead of serially
         self.update_finished: Task = asyncio.create_task(self.update())
-        self.update_finished.add_done_callback(self.handle_exception)
+        self.update_finished.add_done_callback(self.handle_task_exception)
 
         await self.update_finished
         self.print(

slips_files/common/abstracts/iasync_module.py

@@ -30,7 +30,7 @@ def create_task(self, func, *args) -> Task:
         exceptions. because asyncio Tasks do not raise exceptions
         """
         task = asyncio.create_task(func(*args))
-        task.add_done_callback(self.handle_exception)
+        task.add_done_callback(self.handle_task_exception)
 
         # Allow the event loop to run the scheduled task
         # await asyncio.sleep(0)
@@ -39,20 +39,13 @@ def create_task(self, func, *args) -> Task:
         self.tasks.append(task)
         return task
 
-    def handle_exception(self, task):
-        """
-        in asyncmodules we use Async.Task to run some of the functions
-        If an exception occurs in a coroutine that was wrapped in a Task
-        (e.g., asyncio.create_task), the exception does not crash the program
-         but remains in the task.
-        This function is used to handle the exception in the task
-        """
+    def handle_task_exception(self, task: asyncio.Task):
         try:
-            # Access task result to raise the exception if it occurred
-            task.result()
-        except (KeyboardInterrupt, asyncio.exceptions.CancelledError):
-            pass
-        except Exception:
+            exception = task.exception()
+        except asyncio.CancelledError:
+            return  # Task was cancelled, not an error
+        if exception:
+            self.print(f"Unhandled exception in task: {exception}")
             self.print_traceback()
 
     async def main(self): ...
@@ -74,28 +67,45 @@ def run_async_function(self, func: Callable):
         Returns the Future’s result or raise its exception.
         """
         loop = asyncio.get_event_loop()
-        loop.set_exception_handler(self.handle_exception)
+        loop.set_exception_handler(self.handle_loop_exception)
         return loop.run_until_complete(func())
 
+    def handle_loop_exception(self, loop, context):
+        """A common loop exception handler"""
+        exception = context.get("exception")
+        future = context.get("future")
+
+        if future:
+            try:
+                future.result()
+            except Exception:
+                self.print_traceback()
+        elif exception:
+            self.print(f"Unhandled loop exception: {exception}")
+        else:
+            self.print(f"Unhandled loop error: {context.get('message')}")
+
     def run(self):
-        loop = asyncio.new_event_loop()
-        asyncio.set_event_loop(loop)
+        asyncio.run(self._run_pre_main_and_main())
+
+    async def _run_pre_main_and_main(self):
+        """
+        runs pre_main() once, then runs main() in a loop
+        """
+        loop = asyncio.get_event_loop()
+        loop.set_exception_handler(self.handle_loop_exception)
 
         try:
             error: bool = self.pre_main()
             if error or self.should_stop():
-                self.run_async_function(
-                    self.gather_tasks_and_shutdown_gracefully
-                )
+                await self.gather_tasks_and_shutdown_gracefully()
                 return
         except KeyboardInterrupt:
-            self.run_async_function(self.gather_tasks_and_shutdown_gracefully)
+            await self.gather_tasks_and_shutdown_gracefully()
             return
         except RuntimeError as e:
             if "Event loop stopped before Future completed" in str(e):
-                self.run_async_function(
-                    self.gather_tasks_and_shutdown_gracefully
-                )
+                await self.gather_tasks_and_shutdown_gracefully()
                 return
         except Exception:
             self.print_traceback()
@@ -104,18 +114,14 @@ def run(self):
         while True:
             try:
                 if self.should_stop():
-                    self.run_async_function(
-                        self.gather_tasks_and_shutdown_gracefully
-                    )
+                    await self.gather_tasks_and_shutdown_gracefully()
                     return
 
                 # if a module's main() returns 1, it means there's an
                 # error and it needs to stop immediately
-                error: bool = self.run_async_function(self.main)
+                error: bool | None = await self.main()
                 if error:
-                    self.run_async_function(
-                        self.gather_tasks_and_shutdown_gracefully
-                    )
+                    await self.gather_tasks_and_shutdown_gracefully()
                     return
 
             except KeyboardInterrupt:
@@ -128,9 +134,7 @@ def run(self):
                 continue
             except RuntimeError as e:
                 if "Event loop stopped before Future completed" in str(e):
-                    self.run_async_function(
-                        self.gather_tasks_and_shutdown_gracefully
-                    )
+                    await self.gather_tasks_and_shutdown_gracefully()
                     return
             except Exception:
                 self.print_traceback()

slips_files/core/database/redis_db/alert_handler.py

@@ -438,45 +438,6 @@ def update_max_threat_level(
 
         return old_max_threat_level_float
 
-    def update_past_threat_levels(self, profileid, threat_level, confidence):
-        """
-        updates the past_threat_levels key of the given profileid
-        if the past threat level and confidence
-        are the same as the ones we wanna store, we replace the timestamp only
-        """
-        now = utils.convert_ts_format(time.time(), utils.alerts_format)
-        confidence = f"confidence: {confidence}"
-        # this is what we'll be storing in the db, tl, ts, and confidence
-        threat_level_data = (threat_level, now, confidence)
-
-        past_threat_levels: str = self.r.hget(profileid, "past_threat_levels")
-        if past_threat_levels:
-            # get the list of ts and past threat levels
-            past_threat_levels: List[Tuple] = json.loads(past_threat_levels)
-
-            latest: Tuple = past_threat_levels[-1]
-            latest_threat_level: str = latest[0]
-            latest_confidence: str = latest[2]
-
-            if (
-                latest_threat_level == threat_level
-                and latest_confidence == confidence
-            ):
-                # if the past threat level and confidence
-                # are the same as the ones we wanna store,
-                # replace the timestamp only
-                past_threat_levels[-1] = threat_level_data
-                # dont change the old max tl
-            else:
-                # add this threat level to the list of past threat levels
-                past_threat_levels.append(threat_level_data)
-        else:
-            # first time setting a threat level for this profile
-            past_threat_levels = [threat_level_data]
-
-        past_threat_levels = json.dumps(past_threat_levels)
-        self.r.hset(profileid, "past_threat_levels", past_threat_levels)
-
     def update_ips_info(self, profileid, max_threat_lvl, confidence):
         """
         sets the score and confidence of the given ip in the db
@@ -507,11 +468,8 @@ def update_threat_level(
         Do not call this function directy from the db, always call it user
         dbmanager.update_threat_level() to update the trustdb too:D
         """
-
         self.r.hset(profileid, "threat_level", threat_level)
 
-        self.update_past_threat_levels(profileid, threat_level, confidence)
-
         max_threat_lvl: float = self.update_max_threat_level(
             profileid, threat_level
         )