Fixed major issues

dump_memory.py: Added failsafe check for disk space

Logger.py: Made the @function decorator better, by having more precision in its output

vulnscan.py: Added a threadlock as a failsafe to not overload the system

Signed-off-by: Shahm Najeeb <Nirt_12023@outlook.com>

Author: DefinetlyNotAI

CODE/dump_memory.py

@@ -62,11 +62,20 @@ def gather_system_info():
 def memory_dump():
     log.info("Creating basic memory dump scan...")
     pid = os.getpid()
+
     try:
         process = psutil.Process(pid)
         with open("Ram_Dump.txt", "wb") as dump_file:
             total_size = 0
             for mem_region in process.memory_maps(grouped=False):
+                # Check available disk space
+                if os.path.exists("Ram_Dump.txt"):
+                    required_space = LIMIT_FILE_SIZE * 1024 * 1024 * 1.5  # 2x safety margin
+                    free_space = psutil.disk_usage(".").free
+                    if free_space < required_space:
+                        log.error(f"Not enough disk space. Need {required_space / 1024 / 1024:.2f}MB")
+                    return
+
                 # Check if the memory region is readable ('r' permission)
                 if 'r' in mem_region.perms:
                     # Extract start and end addresses from the memory region string
@@ -112,6 +121,7 @@ def memory_dump():
                         total_size += len(metadata_bytes)
                     except Exception as e:
                         log.error(f"Error writing memory region metadata: {str(e)}")
+
     except psutil.Error as e:
         log.error(f"Error opening process memory: {str(e)}")
     except Exception as e:

CODE/logicytics/Logger.py

@@ -273,6 +273,8 @@ def function(self, func: callable):
         :param func: The function to be decorated.
         :return: The wrapper function.
         """
+        if not callable(func):
+            self.exception(f"Function {func.__name__} is not callable.", TypeError)
 
         def wrapper(*args, **kwargs):
             """
@@ -283,14 +285,14 @@ def wrapper(*args, **kwargs):
             :raises TypeError: If the provided function is not callable.
             :return: The result of the decorated function.
             """
-            if not callable(func):
-                self.exception(f"Function {func.__name__} is not callable.", TypeError)
             start_time = time.perf_counter()
-            self.debug(f"Running the function {func.__name__}().")
+            func_args = ", ".join([str(arg) for arg in args] +
+                                  [f"{k}={v}" for k, v in kwargs.items()])
+            self.debug(f"Running the function {func.__name__}({func_args}).")
             result = func(*args, **kwargs)
             end_time = time.perf_counter()
             elapsed_time = end_time - start_time
-            self.debug(f"Function {func.__name__}() executed in {elapsed_time:.6f} seconds.")
+            self.debug(f"{func.__name__}({func_args}) executed in {elapsed_time} -> returned {type(result).__name__}")
             return result
 
         return wrapper

CODE/packet_sniffer.py

@@ -197,7 +197,7 @@ def save_packet_data_to_csv(file_path: str):
         log.warning("No packet data to save.")
 
 
-# Function to visualize the graph
+# Function to visualize the graph of packet connections
 def visualize_graph(node_colors: str = None, node_sizes: str = None):
     """
     Visualizes the graph of packet connections with customizable node colors and sizes.

CODE/sensitive_data_miner.py

@@ -88,7 +88,7 @@ def __search_and_copy_files(cls, keyword: str):
             destination.mkdir()
 
         with ThreadPoolExecutor() as executor:
-            for root, dirs, files in os.walk(drives_root):
+            for root, dirs, _ in os.walk(drives_root):
                 future_to_file = {
                     executor.submit(cls.__search_files_by_keyword, Path(root) / sub_dir, keyword): sub_dir
                     for sub_dir in dirs
@@ -107,7 +107,8 @@ def passwords(cls):
         Returns:
             None
         """
-        keywords = ["password", "secret", "code", "login", "api", "key"]
+        keywords = ["password", "secret", "code", "login", "api", "key",
+                    "token", "auth", "credentials", "private", ]
 
         # Ensure the destination directory is clean
         destination = Path("Password_Files")

CODE/vulnscan.py

@@ -4,6 +4,7 @@
 import os
 import threading
 import warnings
+from concurrent.futures import ThreadPoolExecutor
 
 import joblib
 import numpy as np
@@ -67,8 +68,8 @@ def scan_path(model_path: str, scan_paths: str, vectorizer_path: str):
                 log.info(f"Loading vectorizer from {vectorizer_path}")
                 vectorizer_to_use = joblib.load(vectorizer_path)
         vulnscan(model_to_use, scan_paths, vectorizer_to_use)
-    except Exception as e:
-        log.error(f"Error scanning path {scan_paths}: {e}")
+    except Exception as err:
+        log.error(f"Error scanning path {scan_paths}: {err}")
 
 
 def is_sensitive(model: torch.nn.Module, vectorizer: TfidfVectorizer, file_content: str) -> tuple[bool, float, str]:
@@ -177,12 +178,20 @@ def vulnscan(model, SCAN_PATH, vectorizer):
     # Start scanning
     log.warning("Starting scan - This may take hours and consume memory!!")
 
-    for path in paths:
-        thread = threading.Thread(target=scan_path,
-                                  args=("VulnScan/Model SenseMini .3n3.pth",
-                                        path, "VulnScan/Vectorizer .3n3.pkl"))
-        threads.append(thread)
-        thread.start()
-
-    for thread in threads:
-        thread.join()
+    # Use max_workers based on CPU count but cap it at a reasonable number
+    max_workers = min(32, os.cpu_count() * 2)
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        futures = [
+            executor.submit(
+                scan_path,
+                "VulnScan/Model SenseMini .3n3.pth",
+                path,
+                "VulnScan/Vectorizer .3n3.pkl"
+            )
+            for path in paths
+        ]
+        for future in futures:
+            try:
+                future.result()
+            except Exception as e:
+                log.error(f"Scan failed: {e}")