Made many refactoring changes, fixed bugs and improved code

[__init__.py]
    Added new ObjectLoadError Exception
[_dev.py]
    Improved usage of color_print, as well as added a max attempt of 10 for inputting a valid version
[Flag.py]
    Fixed potential bug where replacing "--" can result in incorrect flags, now its lstrip('-'),
    Fixed bug where if a flag match was found (using cosine sim) it would still say "Sorry" and load up the ML to get the flag for you...
[Logicytics.py]
    Fixed bug where MAX_WORKERS was str rather than int due to improper config.get call
    Added safe_remove and safe_append to make sure the execution list is proper and doesn't break
    Refactored generate_execution_list to be more efficient
    Performance flag now logs memory usage!!!
    Added simple KeyboardInterrupt warning
[network_psutil.py]
    Made the measure method an async to improve efficiency and remove the time.sleep() bottleneck
[packet_sniffer.py]
    Made global variables inside the Sniff class (under __init__)
[vulnscan.py]
    Added checks for vectorizer loading to dodge crashes
    Improved IO bottleneck by making the write happen once in the end rather than every for-loop

Signed-off-by: Shahm Najeeb <[email protected]>

Author: DefinetlyNotAI

.idea/csv-editor.xml

@@ -7,8 +7,8 @@
 import zipfile
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from datetime import datetime
-from typing import Any
 
+import psutil
 from prettytable import PrettyTable
 
 from logicytics import Log, Execute, Check, Get, FileManagement, Flag, DEBUG, DELETE_LOGS
@@ -18,16 +18,33 @@
 ACTION, SUB_ACTION = None, None
 config = configparser.ConfigParser()
 config.read("config.ini")
-MAX_WORKERS = config.get("Settings", "max_workers", fallback=None)
+MAX_WORKERS = config.getint("Settings", "max_workers", fallback=min(32, (os.cpu_count() or 1) + 4))
 log.debug(f"MAX_WORKERS: {MAX_WORKERS}")
 
 
 class ExecuteScript:
     def __init__(self):
-        self.execution_list = self.generate_execution_list()
+        self.execution_list = self.__generate_execution_list()
 
     @staticmethod
-    def generate_execution_list() -> list | list[str] | list[str | Any]:
+    def __safe_remove(file_name: str, file_list: list[str] | set[str]) -> list[str]:
+        file_set = set(file_list)
+        if file_name in file_set:
+            file_set.remove(file_name)
+        else:
+            log.critical(f"The file {file_name} should exist in this directory - But was not found!")
+        return list(file_set)
+
+    @staticmethod
+    def __safe_append(file_name: str, file_list: list[str] | set[str]) -> list[str]:
+        file_set = set(file_list)
+        if os.path.exists(file_name):
+            file_set.add(file_name)
+        else:
+            log.critical(f"Missing required file: {file_name}")
+        return list(file_set)
+
+    def __generate_execution_list(self) -> list[str]:
         """
         Generate an execution list of scripts based on the specified action.
 
@@ -51,11 +68,14 @@ def generate_execution_list() -> list | list[str] | list[str | Any]:
             - Warns users about potential long execution times for certain actions
         """
         execution_list = Get.list_of_files(".", extensions=(".py", ".exe", ".ps1", ".bat"))
-        execution_list.remove("sensitive_data_miner.py")
-        execution_list.remove("dir_list.py")
-        execution_list.remove("tree.ps1")
-        execution_list.remove("vulnscan.py")
-        execution_list.remove("event_log.py")
+        files_to_remove = {
+            "sensitive_data_miner.py",
+            "dir_list.py",
+            "tree.ps1",
+            "vulnscan.py",
+            "event_log.py",
+        }
+        execution_list = [file for file in execution_list if file not in files_to_remove]
 
         if ACTION == "minimal":
             execution_list = [
@@ -69,7 +89,7 @@ def generate_execution_list() -> list | list[str] | list[str | Any]:
                 "event_log.py",
             ]
 
-        if ACTION == "nopy":
+        elif ACTION == "nopy":
             execution_list = [
                 "browser_miner.ps1",
                 "netadapter.ps1",
@@ -78,24 +98,36 @@ def generate_execution_list() -> list | list[str] | list[str | Any]:
                 "tree.ps1"
             ]
 
-        if ACTION == "modded":
+        elif ACTION == "modded":
             # Add all files in MODS to execution list
             execution_list = Get.list_of_files("../MODS",
                                                extensions=(".py", ".exe", ".ps1", ".bat"),
                                                append_file_list=execution_list)
 
-        if ACTION == "depth":
+        elif ACTION == "depth":
             log.warning(
                 "This flag will use clunky and huge scripts, and so may take a long time, but reap great rewards.")
-            execution_list.append("sensitive_data_miner.py")
-            execution_list.append("dir_list.py")
-            execution_list.append("tree.ps1")
-            execution_list.append("event_log.py")
+            files_to_append = {
+                "sensitive_data_miner.py",
+                "dir_list.py",
+                "tree.ps1",
+                "event_log.py",
+            }
+            for file in files_to_append:
+                execution_list = self.__safe_append(file, execution_list)
             log.warning("This flag will use threading!")
 
-        if ACTION == "vulnscan_ai":
+        elif ACTION == "vulnscan_ai":
             # Only vulnscan detector
-            execution_list = ["vulnscan.py"]
+            if os.path.exists("vulnscan.py"):
+                execution_list = ["vulnscan.py"]
+            else:
+                log.critical("Vulnscan is missing...")
+                exit(1)
+
+        if len(execution_list) == 0:
+            log.critical("Nothing is in the execution list.. This is due to faulty code or corrupted Logicytics files!")
+            exit(1)
 
         log.debug(f"The following will be executed: {execution_list}")
         return execution_list
@@ -159,27 +191,43 @@ def __default(self):
 
     def __performance(self):
         """Checks performance of each script."""
+        if DEBUG.lower() != "debug":
+            log.warning("Advised to turn on DEBUG logging!!")
+
         execution_times = []
+        memory_usage = []
+        process = psutil.Process()
 
         for file in range(len(self.execution_list)):
             start_time = datetime.now()
+            start_memory = process.memory_info().rss / 1024 / 1024  # MB
             log.execution(Execute.script(self.execution_list[file]))
             end_time = datetime.now()
+            end_memory = process.memory_info().rss / 1024 / 1024  # MB
             elapsed_time = end_time - start_time
+            memory_delta = end_memory - start_memory
+            memory_usage.append((self.execution_list[file], str(memory_delta)))
             execution_times.append((self.execution_list[file], elapsed_time))
             log.info(f"{self.execution_list[file]} executed in {elapsed_time}")
+            log.info(f"{self.execution_list[file]} used {memory_delta:.2f}MB of memory")
+            log.debug(f"Started with {start_memory}MB of memory and ended with {end_memory}MB of memory")
 
         table = PrettyTable()
-        table.field_names = ["Script", "Execution Time"]
+        table.field_names = ["Script", "Execution Time", "Memory Usage (MB)"]
         for script, elapsed_time in execution_times:
-            table.add_row([script, elapsed_time])
+            memory = next(m[1] for m in memory_usage if m[0] == script)
+            table.add_row([script, elapsed_time, f"{memory:.2f}"])
 
         try:
             with open(
                     f"../ACCESS/LOGS/PERFORMANCE/Performance_Summary_"
                     f"{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.txt", "w"
             ) as f:
                 f.write(table.get_string())
+                f.write(
+                    "\nSome values may be negative, Reason may be due to external resources playing with memory usage, "
+                    "close background tasks to get more accurate readings")
+                f.write("Note: This is not a low-level memory logger, data here isn't 100% accurate!")
             log.info("Performance check complete! Performance log found in ACCESS/LOGS/PERFORMANCE")
         except Exception as e:
             log.error(f"Error writing performance log: {e}")
@@ -454,7 +502,12 @@ def Logicytics():
 
 
 if __name__ == "__main__":
-    Logicytics()
+    try:
+        Logicytics()
+    except KeyboardInterrupt:
+        log.warning("Shutting down Logicytics utility with force causes many files to remain where they shouldn't")
+        log.warning("Please don't force shut Logicytics again - As we don't have a cleanup function yet.")
+        exit(0)
 else:
     log.error("This script cannot be imported!")
     exit(1)

CODE/Logicytics.py

@@ -125,22 +125,32 @@ def _handle_file_operations() -> None:
         if clean_f not in files and clean_f not in EXCLUDE_FILES:
             removed_files.append(clean_f)
 
-    print("\n".join([f"\033[92m+ {file}\033[0m" for file in added_files]))  # Green +
-    print("\n".join([f"\033[91m- {file}\033[0m" for file in removed_files]))  # Red -
-    print("\n".join([f"* {file}" for file in normal_files]))
+    for file in added_files:
+        color_print(f"+ {file}", "green")
+    for file in removed_files:
+        color_print(f"- {file}", "red")
+    for file in normal_files:
+        print(f"* {file}")
 
     if not _prompt_user("Does the list above include your added files?"):
         color_print("[x] Something went wrong! Please contact support.", "red")
         return
 
+    max_attempts = 10
+    attempts = 0
     _update_ini_file("config.ini", files, "files")
     while True:
         version = input(f"\033[36m[?] Enter the new version of the project (Old version is {VERSION}): \033[0m")
+        if attempts >= max_attempts:
+            color_print("[x] Maximum attempts reached. Please run the script again.", "red")
+            exit()
         if re.match(r'^(0|[1-9]\d*)\.(0|[1-9]\d*)\.(0|[1-9]\d*)$', version):
             _update_ini_file("config.ini", version, "version")
             break
         else:
             color_print("[!] Please enter a valid version number (e.g., 1.2.3)", "yellow")
+            attempts += 1
+            color_print(f"[!] {max_attempts - attempts} attempts remaining", "yellow")
     color_print("\n[-] Great Job! Please tick the box in the GitHub PR request for completing steps in --dev", "green")
 
 

CODE/_dev.py

@@ -286,7 +286,7 @@ def update_history(cls, user_input: str, matched_flag: str, accuracy: float):
         if not SAVE_PREFERENCES:
             return
         history_data = cls.load_history()
-        matched_flag = matched_flag.replace("--", "")
+        matched_flag = matched_flag.lstrip('-')
 
         # Ensure that interactions is a dictionary (not a list)
         if not isinstance(history_data['interactions'], dict):
@@ -653,7 +653,8 @@ def __suggest_flag(cls, user_input: str, valid_flags: list[str]):
         # Get the closest valid flag match based on the user's input
         closest_matches = difflib.get_close_matches(user_input, valid_flags, n=1, cutoff=0.6)
         if closest_matches:
-            print(f"Invalid flag '{user_input}', Did you mean '--{closest_matches[0]}'?")
+            print(f"Invalid flag '{user_input}', Did you mean '--{closest_matches[0].replace('_', '-')}'?")
+            exit(1)
 
         # Prompt the user for a description if no close match is found
         user_input_desc = input("We can't find a match, Please provide a description: ").lower()
@@ -663,9 +664,10 @@ def __suggest_flag(cls, user_input: str, valid_flags: list[str]):
         descriptions_list = [f"Run Logicytics with {flag}" for flag in valid_flags]
         flag_received, accuracy_received = _Match.flag(user_input_desc, flags_list, descriptions_list)
         if DEBUG_MODE:
-            print(f"User input: {user_input_desc}\nMatched flag: {flag_received}\nAccuracy: {accuracy_received:.2f}%\n")
+            print(
+                f"User input: {user_input_desc}\nMatched flag: {flag_received.replace('_', '-')}\nAccuracy: {accuracy_received:.2f}%\n")
         else:
-            print(f"Matched flag: {flag_received} (Accuracy: {accuracy_received:.2f}%)\n")
+            print(f"Matched flag: {flag_received.replace('_', '-')} (Accuracy: {accuracy_received:.2f}%)\n")
 
     @staticmethod
     def show_help_menu(return_output: bool = False) -> str | None:

CODE/logicytics/Flag.py

@@ -17,6 +17,13 @@
 Flag = Flag()
 
 
+class ObjectLoadError(Exception):
+    """Raised when an Object fails to load."""
+
+    def __init__(self, message="Failed to load object"):
+        super().__init__(message)
+
+
 def config_data() -> tuple[str, str, list[str], bool]:
     """
     Retrieves configuration data from the 'config.ini' file.
@@ -86,13 +93,13 @@ def deprecated(removal_version: str, reason: str, show_trace: bool = __show_trac
     def decorator(func: callable) -> callable:
         """
         Decorator function that marks a function as deprecated and provides a warning when the function is called.
-        
+
         Args:
             func (callable): The function to be decorated with a deprecation warning.
-        
+
         Returns:
             callable: A wrapper function that preserves the original function's metadata and prints a deprecation warning.
-        
+
         Notes:
             - Uses functools.wraps to preserve the original function's metadata
             - Prints a colorized deprecation warning to stderr
@@ -103,14 +110,14 @@ def decorator(func: callable) -> callable:
         def wrapper(*args, **kwargs) -> callable:
             """
             Wraps a deprecated function to print a warning message before execution.
-            
+
             Args:
                 *args: Positional arguments passed to the original function.
                 **kwargs: Keyword arguments passed to the original function.
-            
+
             Returns:
                 The return value of the original function after printing a deprecation warning.
-            
+
             Warns:
                 Prints a colored deprecation warning to stderr with details about:
                 - Function name being deprecated

CODE/logicytics/__init__.py

@@ -1,6 +1,6 @@
+import asyncio
 import os
 import socket
-import time
 
 import psutil
 
@@ -125,7 +125,7 @@ def __fetch_network_connections_with_process_info(self):
         self.__save_data("network_connections_with_processes.txt", connections_data)
         log.info("Network connections with process info saved.")
 
-    def __measure_network_bandwidth_usage(self, sample_count: int = 5, interval: float = 1.0):
+    async def __measure_network_bandwidth_usage(self, sample_count: int = 5, interval: float = 1.0):
         """
         Measures and saves the average network bandwidth usage.
 
@@ -138,7 +138,7 @@ def __measure_network_bandwidth_usage(self, sample_count: int = 5, interval: flo
         samples = []
         for _ in range(sample_count):
             net1 = psutil.net_io_counters()
-            time.sleep(interval)
+            await asyncio.sleep(interval)
             net2 = psutil.net_io_counters()
             samples.append({
                 'up': (net2.bytes_sent - net1.bytes_sent) / 1024,