atomicshop 3.2.13__py3-none-any.whl → 3.3.1__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__ = '3.2.13'
+__version__ = '3.3.1'

atomicshop/basics/multiprocesses.py

@@ -7,10 +7,77 @@ import concurrent.futures
 from concurrent.futures import ProcessPoolExecutor, as_completed
 from collections import deque
 from typing import Callable
+import time
 
 from ..import system_resources
 
 
+def kill_processes(
+        processes: list
+):
+    """Terminate all children with SIGTERM (or SIGKILL if you like)."""
+    # Ask OS to terminate all processes in the list.
+    for p in processes:
+        if p.is_alive():
+            p.terminate()
+    time.sleep(1)  # give processes a chance to exit cleanly
+    # Force kill all processes in the list.
+    for p in processes:
+        if p.is_alive():
+            p.kill()
+    for p in processes:          # wait for everything to disappear
+        p.join()
+
+
+def is_process_crashed(
+        processes: list[multiprocessing.Process]
+) -> tuple[int, str] | tuple[None, None]:
+    """
+    Check if any of the processes in the list is not alive.
+    :param processes: list, list of multiprocessing.Process objects.
+    :return: tuple(int, string) or None.
+        tuple(0 if any finished cleanly, process name).
+        tuple(1 (or exit code integer) if any process crashed, process_name).
+        None if all processes are still alive.
+
+    ==============================================
+
+    Usage example:
+    processes = [multiprocessing.Process(target=some_function) for _ in range(5)]
+
+    for p in processes:
+        p.start()
+
+    # Check if any process has crashed
+    try:
+        while True:
+            # Poll every second; you can use a shorter sleep if you prefer.
+            result, process_name = is_process_crashed(processes)
+            # If result is None, all processes are still alive.
+            if result is not None:
+                # If result is 0 or 1, we can exit the loop.
+                print(f"Process [{process_name}] finished with exit code {result}.")
+                break
+            time.sleep(1)
+    except KeyboardInterrupt:
+        print("Ctrl-C caught – terminating children…")
+        kill_all(processes)
+        sys.exit(0)
+    """
+
+    for p in processes:
+        if p.exitcode is not None:  # the process is *dead*
+            kill_processes(processes)  # stop the rest
+            if p.exitcode == 0:
+                # print(f"{p.name} exited cleanly; shutting down.")
+                return 0, p.name
+            else:
+                # print(f"{p.name} crashed (exitcode {p.exitcode}). Shutting everything down.")
+                return p.exitcode, p.name
+
+    return None, None  # all processes are still alive
+
+
 def process_wrap_queue(function_reference: Callable, *args, **kwargs):
     """
     The function receives function reference and arguments, and executes the function in a thread.

atomicshop/config_init.py

@@ -8,7 +8,11 @@ CONFIG_FILE_NAME = 'config.toml'
 CONFIG: dict = dict()
 
 
-def get_config(script_directory: str = None, config_file_name: str = CONFIG_FILE_NAME) -> dict:
+def get_config(
+        script_directory: str = None,
+        config_file_name: str = CONFIG_FILE_NAME,
+        print_kwargs: dict = None
+) -> dict:
     """
     Get the config file content.
 
@@ -16,6 +20,7 @@ def get_config(script_directory: str = None, config_file_name: str = CONFIG_FILE
         get the working directory instead.
     :param config_file_name: string, name of the config file. Default is 'config.toml' as specified in the constant:
         'CONFIG_FILE_NAME'.
+    :param print_kwargs: dict, additional arguments to pass to the print function. Default is None.
     :return: dict.
     """
 
@@ -25,7 +30,7 @@ def get_config(script_directory: str = None, config_file_name: str = CONFIG_FILE
     if not script_directory:
         script_directory = filesystem.get_working_directory()
 
-    CONFIG = tomls.read_toml_file(f'{script_directory}{os.sep}{config_file_name}')
+    CONFIG = tomls.read_toml_file(f'{script_directory}{os.sep}{config_file_name}', **(print_kwargs or {}))
     return CONFIG
 
 

atomicshop/mitm/config_static.py

@@ -161,7 +161,10 @@ class ProcessName:
     ssh_script_to_execute: Literal['process_from_port', 'process_from_ipv4'] = 'process_from_port'
 
 
-def load_config(config_toml_file_path: str):
+def load_config(
+        config_toml_file_path: str,
+        print_kwargs: dict = None
+):
     # global CONFIG
 
     script_path = os.path.dirname(config_toml_file_path)
@@ -169,11 +172,11 @@ def load_config(config_toml_file_path: str):
     MainConfig.update()
 
     # Load the configuration file.
-    result = import_config.import_config_files(config_toml_file_path)
+    result = import_config.import_config_files(config_toml_file_path, print_kwargs=print_kwargs or {})
     return result
 
 
-def get_listening_addresses(client_message: ClientMessage) -> dict:
+def get_listening_addresses(client_message: ClientMessage) -> dict | None:
     """
     Get the list of listening addresses from the TCPServer configuration.
     If no_engines_usage_to_listen_addresses_enable is True, return the no_engines_listening_address_list.

atomicshop/mitm/connection_thread_worker.py

@@ -10,20 +10,25 @@ from ..basics import threads, tracebacks
 from ..print_api import print_api
 
 from .message import ClientMessage
-from . import config_static, initialize_engines
+from . import initialize_engines
+from ..wrappers.loggingw import loggingw
+# This is needed only for the data typing.
+from . import config_static as cf
 
 
 def thread_worker_main(
+        # These parameters come from the SocketWrapper.
         client_socket,
         process_commandline: str,
         is_tls: bool,
         tls_type: str,
         tls_version: str,
         domain_from_dns,
-        network_logger,
         statistics_writer,
         engines_list,
-        reference_module
+
+        # These parameters come from the main mitm module.
+        config_static: cf
 ):
     def output_statistics_csv_row(client_message: ClientMessage):
         # If there is no '.code' attribute in HTTPResponse, this means that this is not an HTTP message, so there is no
@@ -543,6 +548,8 @@ def thread_worker_main(
     # ================================================================================================================
     # This is the start of the thread_worker_main function
 
+    network_logger = loggingw.get_logger_with_level(config_static.MainConfig.LOGGER_NAME)
+
     # Only protocols that are encrypted with TLS have the server name attribute.
     if is_tls:
         # Get current destination domain
@@ -568,7 +575,7 @@ def thread_worker_main(
     found_domain_module = initialize_engines.assign_class_by_domain(
         engines_list=engines_list,
         message_domain_name=server_name,
-        reference_module=reference_module
+        reference_module=config_static.REFERENCE_MODULE
     )
     parser = found_domain_module.parser_class_object
     requester = found_domain_module.requester_class_object()

atomicshop/mitm/engines/__reference_general/responder___reference_general.py

@@ -212,3 +212,24 @@ class ResponderGeneral(ResponderParent):
     #
     #     result_response_list: list[bytes] = [byte_response]
     #     return result_response_list
+    #
+    # ==================================================================================================================
+    # TEST RESPONSE.
+    # def create_response(self, class_client_message: ClientMessage):
+    #     resp_body_text: bytes = b"<html><body>TEST OK!</body></html>\n"
+    #     resp_status_code: int = 200
+    #     resp_headers: dict = {
+    #         # Tell the browser it’s plain text (could be “text/html” if you wrap it in HTML).
+    #         "Content-Type": "text/html; charset=utf-8"}
+    #
+    #     # Build the raw bytes to send.
+    #     byte_response = self.build_byte_response(
+    #         http_version="HTTP/1.1",
+    #         status_code=resp_status_code,
+    #         headers=resp_headers,
+    #         body=resp_body_text
+    #
+    #     )
+    #
+    #     result_response_list: list[bytes] = [byte_response]
+    #     return result_response_list

atomicshop/mitm/import_config.py

@@ -35,17 +35,22 @@ def assign_bool(dict_instance: dict, section: str, key: str):
 
 
 def import_config_files(
-        config_file_path: str
+        config_file_path: str,
+        print_kwargs: dict = None
 ):
     """
     Import the configuration file 'config.toml' and write all the values to 'config_static' dataclasses module.
 
     :param config_file_path:
+    :param print_kwargs: dict, additional arguments to pass to the print function.
     :return:
     """
 
     config_toml: dict = config_init.get_config(
-        script_directory=str(Path(config_file_path).parent), config_file_name=Path(config_file_path).name)
+        script_directory=str(Path(config_file_path).parent),
+        config_file_name=Path(config_file_path).name,
+        print_kwargs=print_kwargs or {}
+    )
 
     # Assign boolean values to the toml dict module.
     for boolean_tuple in config_static.LIST_OF_BOOLEANS:
@@ -61,7 +66,7 @@ def import_config_files(
 
     manipulations_after_import()
 
-    result = import_engines_configs()
+    result = import_engines_configs(print_kwargs=print_kwargs or {})
     if result != 0:
         return result
 
@@ -69,7 +74,7 @@ def import_config_files(
     return result
 
 
-def import_engines_configs() -> int:
+def import_engines_configs(print_kwargs: dict) -> int:
     """
     Import the engines configuration files and write all the values to 'config_static' dataclasses module.
 
@@ -88,8 +93,8 @@ def import_engines_configs() -> int:
     for engine_config_path in engine_config_path_list:
         # Initialize engine.
         current_module: initialize_engines.ModuleCategory = initialize_engines.ModuleCategory(config_static.MainConfig.SCRIPT_DIRECTORY)
-        current_module.fill_engine_fields_from_config(engine_config_path.path)
-        result_code, error = current_module.initialize_engine()
+        current_module.fill_engine_fields_from_config(engine_config_path.path, print_kwargs=print_kwargs or {})
+        result_code, error = current_module.initialize_engine(print_kwargs=print_kwargs or {})
         if result_code != 0:
             print_api(f"Error initializing engine from directory: {Path(engine_config_path.path).parent}\n{error}", color='red')
             return result_code
@@ -344,6 +349,7 @@ def manipulations_after_import():
         config_static.Certificates.custom_private_key_path = filesystem.check_absolute_path___add_full(
             config_static.Certificates.custom_private_key_path, config_static.MainConfig.SCRIPT_DIRECTORY)
     else:
+        # noinspection PyTypeChecker
         config_static.Certificates.custom_private_key_path = None
 
     config_static.Certificates.sni_server_certificates_cache_directory = filesystem.check_absolute_path___add_full(

atomicshop/mitm/initialize_engines.py

@@ -41,9 +41,13 @@ class ModuleCategory:
         self.responder_file_path = reference_folder_path + os.sep + "responder___reference_general.py"
         self.recorder_file_path = reference_folder_path + os.sep + "recorder___reference_general.py"
 
-    def fill_engine_fields_from_config(self, engine_config_file_path: str):
+    def fill_engine_fields_from_config(
+            self,
+            engine_config_file_path: str,
+            print_kwargs: dict = None
+    ):
         # Read the configuration file of the engine.
-        configuration_data = tomls.read_toml_file(engine_config_file_path)
+        configuration_data = tomls.read_toml_file(engine_config_file_path, **(print_kwargs or {}))
 
         engine_directory_path: str = str(Path(engine_config_file_path).parent)
         self.engine_name = Path(engine_directory_path).name
@@ -92,17 +96,21 @@ class ModuleCategory:
         for subdomain, file_name in self.mtls.items():
             self.mtls[subdomain] = f'{engine_directory_path}{os.sep}{file_name}'
 
-    def initialize_engine(self, reference_general: bool = False) -> tuple[int, str]:
+    def initialize_engine(
+            self,
+            reference_general: bool = False,
+            print_kwargs: dict = None
+    ) -> tuple[int, str]:
         try:
             if not reference_general:
                 self.parser_class_object = import_first_class_name_from_file_path(
-                    self.script_directory, self.parser_file_path)
+                    self.script_directory, self.parser_file_path, **(print_kwargs or {}))
                 self.requester_class_object = import_first_class_name_from_file_path(
-                    self.script_directory, self.requester_file_path)
+                    self.script_directory, self.requester_file_path, **(print_kwargs or {}))
                 self.responder_class_object = import_first_class_name_from_file_path(
-                    self.script_directory, self.responder_file_path)
+                    self.script_directory, self.responder_file_path, **(print_kwargs or {}))
                 self.recorder_class_object = import_first_class_name_from_file_path(
-                    self.script_directory, self.recorder_file_path)
+                    self.script_directory, self.recorder_file_path, **(print_kwargs or {}))
             else:
                 self.parser_class_object = parser___reference_general.ParserGeneral
                 self.requester_class_object = requester___reference_general.RequesterGeneral