Merge branch 'bug-fixes-and-v3.3.0' into coderabbitai/docstrings/a11ce93

Author: DefinetlyNotAI

.idea/csv-editor.xml

@@ -122,11 +122,20 @@ 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
@@ -172,6 +181,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/dump_memory.py

@@ -349,6 +349,8 @@ def example_function():
                 # Function implementation
                 pass
         """
+        if not callable(func):
+            self.exception(f"Function {func.__name__} is not callable.", TypeError)
 
         def wrapper(*args, **kwargs):
             """
@@ -371,14 +373,14 @@ def wrapper(*args, **kwargs):
                 - Measures and logs the total execution time with microsecond precision
                 - Preserves the original function's return value
             """
-            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/logicytics/Logger.py

@@ -13,19 +13,15 @@
 if __name__ == "__main__":
     log = Log({"log_level": DEBUG})
 
-try:
-    # Read configuration from config.ini
-    config = ConfigParser()
-    config.read('config.ini')
-    config = config['PacketSniffer Settings']
+# Read configuration from config.ini
+config = ConfigParser()
+config.read('config.ini')
+config = config['PacketSniffer Settings']
 
-    # Global configuration
-    conf.verb = 0  # Turn off verbosity for clean output
-    packet_data = []  # List to store packet information
-    G = nx.Graph()  # Initialize a graph
-except Exception as e:
-    log.error(f"Error reading configuration: {e}")
-    exit(1)
+# Global configuration
+conf.verb = 0  # Turn off verbosity for clean output
+packet_data = []  # List to store packet information
+G = nx.Graph()  # Initialize a graph
 
 
 # Function to process and log packet details
@@ -274,7 +270,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.
@@ -345,6 +341,14 @@ def packet_sniffer():
         - Calls start_sniffing() to capture network packets
         - Exits the program if critical configuration errors are encountered
     """
+
+    def correct_interface_name(interface_name: str) -> str:
+        corrections = {
+            "WiFi": "Wi-Fi",
+            "Wi-Fi": "WiFi"
+        }
+        return corrections.get(interface_name, interface_name)
+
     interface = config['interface']
     packet_count = int(config['packet_count'])
     timeout = int(config['timeout'])
@@ -360,18 +364,16 @@ def packet_sniffer():
             log.error("Error reading configuration: Improper values for packet count or timeout")
         exit(1)
 
-    try:
-        start_sniffing(interface, packet_count, timeout)
-    except Exception as err:
-        log.error(f"Invalid interface '{interface}'. Please check the configuration: {err}")
-        if interface == "WiFi" or interface == "Wi-Fi":
-            log.warning("Attempting to correct the interface name...")
-            interface = "Wi-Fi" if interface == "WiFi" else "WiFi"
-            log.info(f"Interface name auto-corrected to '{interface}', retrying packet sniffing...")
-            try:
-                start_sniffing(interface, packet_count, timeout)
-            except Exception as err:
-                log.error(f"Error sniffing packets on auto-corrected interface '{interface}': {err}")
+    for attempt in range(2):  # Try original and corrected name
+        try:
+            start_sniffing(interface, packet_count, timeout)
+            break
+        except Exception as err:
+            if attempt == 0 and interface in ("WiFi", "Wi-Fi"):
+                log.warning(f"Retrying with corrected interface name...")
+                interface = correct_interface_name(interface)
+            else:
+                log.error(f"Failed to sniff packets: {err}")
 
 
 # Entry point of the script
@@ -380,4 +382,6 @@ def packet_sniffer():
         packet_sniffer()
     except Exception as e:
         log.error(e)
-        exit(1)
+    finally:
+        if G:
+            plt.close()

CODE/packet_sniffer.py

@@ -120,7 +120,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
@@ -159,7 +159,8 @@ def passwords(cls):
         Logging:
             - Logs an informational message upon completion of the search and copy process
         """
-        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/sensitive_data_miner.py

@@ -4,6 +4,7 @@
 import os
 import threading
 import warnings
+from concurrent.futures import ThreadPoolExecutor
 
 import joblib
 import numpy as np
@@ -82,8 +83,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]:
@@ -203,12 +204,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}")