atomicshop 2.21.1__py3-none-any.whl → 3.0.0__py3-none-any.whl

atomicshop/__init__.py

@@ -1,4 +1,4 @@
 """Atomic Basic functions and classes to make developer life easier"""
 
 __author__ = "Den Kras"
-__version__ = '2.21.1'
+__version__ = '3.0.0'

atomicshop/basics/multiprocesses.py

@@ -1,14 +1,17 @@
 import multiprocessing
 import multiprocessing.managers
+import os
 import queue
+import threading
 import concurrent.futures
 from concurrent.futures import ProcessPoolExecutor, as_completed
+from collections import deque
 from typing import Callable
 
 from ..import system_resources
 
 
-def process_wrap_queue(function_reference, *args, **kwargs):
+def process_wrap_queue(function_reference: Callable, *args, **kwargs):
     """
     The function receives function reference and arguments, and executes the function in a thread.
     "_queue" means that a queue.put() is used to store the result of the function and queue.get() to output it.
@@ -154,38 +157,63 @@ class MultiProcessorRecursive:
         self.async_results: list = []
 
     def run_process(self):
-        while self.input_list:
-            new_input_list = []
-            for item in self.input_list:
-                # Check system resources before processing each item
-                system_resources.wait_for_resource_availability(
-                    cpu_percent_max=self.cpu_percent_max,
-                    memory_percent_max=self.memory_percent_max,
-                    wait_time=self.wait_time,
-                    system_monitor_manager_dict=self.system_monitor_manager_dict)
+        """
+        Start with the items currently in self.input_list, but whenever a task
+        finishes schedule the children it returns *right away*.
+        The loop ends when there are no more outstanding tasks.
+        """
+        # ----------  internal helpers  ----------
+        outstanding = 0  # tasks that have been submitted but not yet finished
+        done_event = threading.Event()  # let the main thread wait until work is over
+
+        def _submit(item):
+            nonlocal outstanding
+            # Wait for resources *before* submitting a new job
+            system_resources.wait_for_resource_availability(
+                cpu_percent_max=self.cpu_percent_max,
+                memory_percent_max=self.memory_percent_max,
+                wait_time=self.wait_time,
+                system_monitor_manager_dict=self.system_monitor_manager_dict
+            )
+            outstanding += 1
+            self.pool.apply_async(
+                self.process_function,
+                (item,),
+                callback=_on_finish,  # called in the main process when result is ready
+                error_callback=_on_error
+            )
 
-                # Process the item
-                async_result = self.pool.apply_async(self.process_function, (item,))
-                self.async_results.append(async_result)
+        def _on_finish(result):
+            """Pool calls this in the parent process thread when a job completes."""
+            nonlocal outstanding
+            outstanding -= 1
 
-            # Reset input_list for next round of processing
-            self.input_list = []
+            # The worker returned a list of new items – submit them immediately
+            if result:
+                for child in result:
+                    _submit(child)
+
+            # If no work left, release the waiter
+            if outstanding == 0:
+                done_event.set()
 
-            # Collect results as they complete
-            for async_result in self.async_results:
-                try:
-                    result = async_result.get()
-                    # Assuming process_function returns a list, extend new_input_list
-                    new_input_list.extend(result)
-                except Exception:
-                    raise
-
-            # Update the input_list for the next iteration
-            self.input_list = new_input_list
-            # Clear the async_results for the next iteration
-            self.async_results.clear()
-
-    def shutdown_pool(self):
+        def _on_error(exc):
+            """Propagate the first exception and stop everything cleanly."""
+            done_event.set()
+            raise exc  # let your code deal with it – you can customise this
+
+        # ----------  kick‑off  ----------
+        # Schedule the items we already have
+        for item in self.input_list:
+            _submit(item)
+
+        # Clear the input list; after this point everything is driven by callbacks
+        self.input_list.clear()
+
+        # Wait until all recursively spawned work is finished
+        done_event.wait()
+
+    def shutdown(self):
         """Shuts down the pool gracefully."""
         if self.pool:
             self.pool.close()  # Stop accepting new tasks
@@ -193,6 +221,176 @@ class MultiProcessorRecursive:
             self.pool = None
 
 
+class _MultiProcessorRecursiveWithProcessPoolExecutor:
+    def __init__(
+            self,
+            process_function: Callable,
+            input_list: list,
+            max_workers: int = None,
+            cpu_percent_max: int = 80,
+            memory_percent_max: int = 80,
+            wait_time: float = 5,
+            system_monitor_manager_dict: multiprocessing.managers.DictProxy = None
+    ):
+        """
+        THIS CLASS USES THE concurrent.futures.ProcessPoolExecutor to achieve parallelism.
+        For some reason I got freezes on exceptions without the exception output after the run_process() method finished
+        and the pool remained open. So, using the MultiProcessorRecursive instead.
+
+        MultiProcessor class. Used to execute functions in parallel. The result of each execution is fed back
+            to the provided function. Making it sort of recursive execution.
+        :param process_function: function, function to execute on the input list.
+        :param input_list: list, list of inputs to process.
+        :param max_workers: integer, number of workers to execute functions in parallel. Default is None, which
+            is the number of CPUs that will be counted automatically by the multiprocessing module.
+        :param cpu_percent_max: integer, maximum CPU percentage. Above that usage, we will wait before starting new
+            execution.
+        :param memory_percent_max: integer, maximum memory percentage. Above that usage, we will wait, before starting
+            new execution.
+        :param wait_time: float, time to wait if the CPU or memory usage is above the maximum percentage.
+        :param system_monitor_manager_dict: multiprocessing.managers.DictProxy, shared manager dict for
+            system monitoring. The object is the output of atomicshop.system_resource_monitor.
+            If you are already running this monitor, you can pass the manager_dict to both the system monitor and this
+            class to share the system resources data.
+            If this is used, the system resources will be checked before starting each new execution from this
+            shared dict instead of performing new checks.
+
+        Usage Examples:
+            def unpack_file(file_path):
+                # Process the file at file_path and unpack it.
+                # Return a list of new file paths that were extracted from the provided path.
+                return [new_file_path1, new_file_path2]  # Example return value
+
+            # List of file paths to process
+            file_paths = ["path1", "path2", "path3"]
+
+            # Note: unpack_file Callable is passed to init without parentheses.
+
+            1. Providing the list directly to process at once:
+                # Initialize the processor.
+                processor = MultiProcessor(
+                    process_function=unpack_file,
+                    input_list=file_paths,
+                    max_workers=4,  # Number of parallel workers
+                    cpu_percent_max=80,  # Max CPU usage percentage
+                    memory_percent_max=80,  # Max memory usage percentage
+                    wait_time=5  # Time to wait if resources are overused
+                )
+
+                # Process the list of files at once.
+                processor.run_process()
+                # Shutdown the pool processes after processing.
+                processor.shutdown_pool()
+
+            2. Processing each file in the list differently then adding to the list of the multiprocessing instance then executing.
+                # Initialize the processor once, before the loop, with empty input_list.
+                processor = MultiProcessor(
+                    process_function=unpack_file,
+                    input_list=[],
+                    max_workers=4,  # Number of parallel workers
+                    cpu_percent_max=80,  # Max CPU usage percentage
+                    memory_percent_max=80,  # Max memory usage percentage
+                    wait_time=5  # Time to wait if resources are overused
+                )
+
+                for file_path in file_paths:
+                    # <Process each file>.
+                    # Add the result to the input_list of the processor.
+                    processor.input_list.append(file_path)
+
+                # Process the list of files at once.
+                processor.run_process()
+                # Shutdown the pool processes after processing.
+                processor.shutdown_pool()
+
+            3. Processing each file in the list separately, since we're using an unpacking function that
+               will create more files, but the context for this operation is different for extraction
+               of each main file inside the list:
+
+                # Initialize the processor once, before the loop, with empty input_list.
+                processor = MultiProcessor(
+                    process_function=unpack_file,
+                    input_list=[],
+                    max_workers=4,  # Number of parallel workers
+                    cpu_percent_max=80,  # Max CPU usage percentage
+                    memory_percent_max=80,  # Max memory usage percentage
+                    wait_time=5  # Time to wait if resources are overused
+                )
+
+                for file_path in file_paths:
+                    # <Process each file>.
+                    # Add the result to the input_list of the processor.
+                    processor.input_list.append(file_path)
+                    # Process the added file path separately.
+                    processor.run_process()
+
+                # Shutdown the pool processes after processing.
+                processor.shutdown_pool()
+        """
+
+        self.process_function: Callable = process_function
+        self.input_list: list = input_list
+        self.cpu_percent_max: int = cpu_percent_max
+        self.memory_percent_max: int = memory_percent_max
+        self.wait_time: float = wait_time
+        self.system_monitor_manager_dict: multiprocessing.managers.DictProxy = system_monitor_manager_dict
+
+        if max_workers is None:
+            max_workers = os.cpu_count()
+        self.max_workers: int = max_workers
+
+        # Create the executor once and reuse it.
+        # noinspection PyTypeChecker
+        self.executor: ProcessPoolExecutor = None
+
+    def _ensure_executor(self):
+        """Create a new pool if we do not have one or if the old one was shut."""
+        if self.executor is None or getattr(self.executor, '_shutdown', False):
+            self.executor = ProcessPoolExecutor(max_workers=self.max_workers)
+
+    def run_process(self):
+        # Make sure we have a live executor
+        self._ensure_executor()
+
+        work_q = deque(self.input_list)  # breadth‑first queue
+        self.input_list.clear()
+        futures = set()
+
+        # helper to submit jobs up to the concurrency limit
+        def _fill():
+            while work_q and len(futures) < self.max_workers:
+                item = work_q.popleft()
+                system_resources.wait_for_resource_availability(
+                    cpu_percent_max=self.cpu_percent_max,
+                    memory_percent_max=self.memory_percent_max,
+                    wait_time=self.wait_time,
+                    system_monitor_manager_dict=self.system_monitor_manager_dict
+                )
+                futures.add(self.executor.submit(self.process_function, item))
+
+        _fill()  # start the first wave
+
+        while futures:
+            for fut in as_completed(futures):
+                futures.remove(fut)  # a slot just freed up
+
+                # propagate worker exceptions immediately
+                children = fut.result()
+
+                # schedule the newly discovered items
+                if children:
+                    work_q.extend(children)
+
+                _fill()  # keep the pool saturated
+                break  # leave the for‑loop so as_completed resets
+
+    def shutdown(self):
+        """Shuts down the executor gracefully."""
+        if self.executor:
+            self.executor.shutdown(wait=True)  # blocks until all tasks complete
+            self.executor = None
+
+
 class ConcurrentProcessorRecursive:
     def __init__(
             self,

atomicshop/dns.py

@@ -174,3 +174,5 @@ def default_dns_gateway_main() -> int:
     elif args.dynamic:
         set_connection_dns_gateway_dynamic(
             connection_name=args.connection_name, use_default_connection=args.connection_default)
+
+    return 0

atomicshop/mitm/config_static.py

@@ -1,5 +1,6 @@
 import os
 from dataclasses import dataclass
+from typing import Literal
 
 from . import import_config
 
@@ -157,7 +158,7 @@ class ProcessName:
     ssh_user: str
     ssh_pass: str
 
-    ssh_script_to_execute = 'process_from_port'
+    ssh_script_to_execute: Literal['process_from_port', 'process_from_ipv4'] = 'process_from_port'
 
 
 def load_config(config_toml_file_path: str):

atomicshop/mitm/engines/create_module_template.py

@@ -30,7 +30,7 @@ class CreateModuleTemplate:
     def __init__(self):
         # === Get input variables. ===
         self.engine_name: str = parse_arguments().engine_name
-        self.domains: list = ['example.com']
+        self.domains: list = ['example.com:443', 'example.org:80']
 
         # New engine's directory.
         self.new_engine_directory: str = ENGINES_DIRECTORY_PATH + os.sep + self.engine_name
@@ -76,15 +76,10 @@ class CreateModuleTemplate:
 
         config_lines_list.append('[engine]')
         config_lines_list.append(f'domains = [{", ".join(domains_with_quotes)}]')
-        config_lines_list.append('dns_target = "127.0.0.1"')
-        config_lines_list.append('tcp_listening_address_list = ["0.0.0.0:443"]\n')
+        config_lines_list.append('localhost = 0\n')
         # config_lines_list.append(f'\n')
         config_lines_list.append('[mtls]')
         config_lines_list.append('# "subdomain.domain.com" = "file_name_in_current_dir.pem"\n')
-        # config_lines_list.append(f'\n')
-        config_lines_list.append('[no_sni]')
-        config_lines_list.append('get_from_dns = 1         # Blocking, the accept function will wait until the domain is received from DNS.')
-        config_lines_list.append('serve_domain_on_address = {0 = [{"example.com" = "127.0.0.2:443"}]}')
 
         config_file_path = self.new_engine_directory + os.sep + CONFIG_FILE_NAME
 

atomicshop/mitm/import_config.py

@@ -151,37 +151,41 @@ def check_configurations() -> int:
         print_api(error_message, color="red")
         return 1
 
+    is_localhost: bool | None = None
     for engine in config_static.ENGINES_LIST:
-        if engine.no_sni.get_from_dns and engine.no_sni.serve_domain_on_address_enable:
-            message = (
-                f"Both [get_from_dns] and [serve_domain_on_address] are enabled in [no_sni] section of the engine.\n"
-                f"Only one Can be True.")
-            print_api(message, color="red")
-            return 1
-        if not engine.no_sni.get_from_dns and not engine.no_sni.serve_domain_on_address_enable:
-            message = (
-                f"Both [get_from_dns] and [serve_domain_on_address] are disabled in [no_sni] section of the engine.\n"
-                f"Only one Can be True.")
-            print_api(message, color="red")
-            return 1
-
-        if engine.no_sni.serve_domain_on_address_enable:
-            # Check if the domains in no_sni are the same as in the engine. They should not be.
-            # Same goes for the address.
-            for domain, address_ip_port in engine.no_sni.serve_domain_on_address_dict.items():
-                if domain in engine.domain_list:
+        for domain_port in engine.domain_list:
+            # Check if the domains has port.
+            if ':' not in domain_port:
+                message = (
+                    f"[*] Domain [{domain_port}] doesn't have a port.\n"
+                    f"Please check your engine configuration file.")
+                print_api(message, color="red")
+                return 1
+            else:
+                # Split the domain and port.
+                domain, port = domain_port.split(':')
+                # Check if the port is a number.
+                if not port.isdigit():
                     message = (
-                        f"[*] No SNI setting: The domain [{domain}] is in the engine domains list [{engine.domain_list}].\n"
-                        f"The point of the no_sni section is to serve specific domains on separate addresses.\n")
+                        f"[*] Port [{port}] is not a number.\n"
+                        f"Please check your engine configuration file.")
                     print_api(message, color="red")
                     return 1
 
-                if address_ip_port in engine.tcp_listening_address_list:
-                    message = (
-                        f"[*] No SNI setting: The address [{address_ip_port}] is in the engine listening interfaces list [{engine.tcp_listening_address_list}].\n"
-                        f"The point of the no_sni section is to serve specific domains on separate addresses.\n")
-                    print_api(message, color="red")
-                    return 1
+        # Check if 'localhost' is set in all the engines, or not.
+        # There can't be mixed engines where local host is set and not set.
+        # It can be all engines will be localhost or none of them.
+        if is_localhost is None:
+            is_localhost = engine.is_localhost
+        else:
+            if is_localhost != engine.is_localhost:
+                message = (
+                    f"[*] Mixed [localhost] setting in the engines found.\n"
+                    f"[*] Some engines are set to [localhost] and some are not.\n"
+                    f"[*] This is not allowed. All engines must be set to [localhost = 1] or All engines must be set to [localhost = 0].\n"
+                    f"Please check your engine configuration files.")
+                print_api(message, color="red")
+                return 1
 
     # Check admin right if on localhost ============================================================================
     # If any of the DNS IP target addresses is localhost loopback, then we need to check if the script

atomicshop/mitm/initialize_engines.py

@@ -7,23 +7,15 @@ from .engines.__reference_general import parser___reference_general, responder__
     recorder___reference_general
 
 
-class NoSNI:
-    def __init__(self):
-        self.get_from_dns: bool = False
-        self.serve_domain_on_address_enable: bool = False
-        self.serve_domain_on_address_dict: dict = dict()
-
-
 class ModuleCategory:
     def __init__(self, script_directory: str):
         self.engine_name: str = str()
         self.script_directory: str = script_directory
 
         self.domain_list: list = list()
-        self.dns_target: str = str()
-        self.tcp_listening_address_list: list = list()
+        self.domain_target_dict: dict = dict()
+        self.is_localhost: bool = bool()
         self.mtls: dict = dict()
-        self.no_sni: NoSNI = NoSNI()
 
         self.parser_file_path: str = str()
         self.responder_file_path: str = str()
@@ -51,24 +43,11 @@ class ModuleCategory:
 
         # Getting the parameters from engine config file
         self.domain_list = configuration_data['engine']['domains']
-        self.dns_target = configuration_data['engine']['dns_target']
-        self.tcp_listening_address_list = configuration_data['engine']['tcp_listening_address_list']
+        self.is_localhost = bool(configuration_data['engine']['localhost'])
 
         if 'mtls' in configuration_data:
             self.mtls = configuration_data['mtls']
 
-        self.no_sni.get_from_dns = bool(configuration_data['no_sni']['get_from_dns'])
-
-        for enable_bool, address_list in configuration_data['no_sni']['serve_domain_on_address'].items():
-            if enable_bool in ['0', '1']:
-                self.no_sni.serve_domain_on_address_enable = bool(int(enable_bool))
-            else:
-                raise ValueError(f"Error: no_sni -> serve_domain_on_address -> key must be 0 or 1.")
-
-            for address in address_list:
-                for domain, address_ip_port in address.items():
-                    self.no_sni.serve_domain_on_address_dict = {domain: address_ip_port}
-
         # If there's module configuration file, but no domains in it, there's no point to continue.
         # Since, each engine is based on domains.
         if not self.domain_list or self.domain_list[0] == '':
@@ -86,6 +65,12 @@ class ModuleCategory:
         self.responder_file_path = f"{engine_directory_path}{os.sep}responder{file_name_suffix}.py"
         self.recorder_file_path = f"{engine_directory_path}{os.sep}recorder{file_name_suffix}.py"
 
+        for domain_index, domain_port_string in enumerate(self.domain_list):
+            # Splitting the domain and port
+            domain, port = domain_port_string.split(':')
+
+            self.domain_target_dict[domain] = {'ip': None, 'port': port}
+
         for subdomain, file_name in self.mtls.items():
             self.mtls[subdomain] = f'{engine_directory_path}{os.sep}{file_name}'