atomicshop 3.3.8__py3-none-any.whl → 3.10.0__py3-none-any.whl

atomicshop/wrappers/socketw/{base.py → socket_base.py}

@@ -1,4 +1,5 @@
 import socket
+import time
 
 
 LOCALHOST_IPV4: str = '127.0.0.1'
@@ -99,4 +100,35 @@ def get_host_name_from_ip_address(ip_address: str) -> str:
     host_name, alias_list, ipaddr_list = socket.gethostbyaddr(ip_address)
     _ = alias_list, ipaddr_list
 
-    return host_name
+    return host_name
+
+
+def wait_for_ip_bindable(
+    ip: str,
+    port: int = 0,
+    timeout: float = 15.0,
+    interval: float = 0.5,
+) -> None:
+    """
+    Wait until a single IP is bindable (or timeout).
+
+    Raises TimeoutError if the IP cannot be bound within 'timeout' seconds.
+    """
+    deadline = time.time() + timeout
+    last_err: OSError | None = None
+
+    while time.time() < deadline:
+        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        try:
+            s.bind((ip, port))
+            s.close()
+            return  # success
+        except OSError as e:
+            last_err = e
+            s.close()
+            time.sleep(interval)
+
+    raise TimeoutError(
+        f"IP {ip} not bindable within {timeout} seconds; "
+        f"last error: {last_err}"
+    )

atomicshop/wrappers/socketw/socket_client.py

@@ -31,7 +31,8 @@ class SocketClient:
             dns_servers_list: list[str] = None,
             logger: logging.Logger = None,
             custom_pem_client_certificate_file_path: str = None,
-            enable_sslkeylogfile_env_to_client_ssl_context: bool = False
+            enable_sslkeylogfile_env_to_client_ssl_context: bool = False,
+            sslkeylog_file_path:str = None
     ):
         """
         If you have a certificate for domain, but not for the IPv4 address, the SSL Socket context can be created for
@@ -68,6 +69,7 @@ class SocketClient:
         self.dns_servers_list = dns_servers_list
         self.custom_pem_client_certificate_file_path: str = custom_pem_client_certificate_file_path
         self.enable_sslkeylogfile_env_to_client_ssl_context: bool = enable_sslkeylogfile_env_to_client_ssl_context
+        self.sslkeylog_file_path: str = sslkeylog_file_path
 
         if logger:
             # Create child logger for the provided logger with the module's name.
@@ -101,7 +103,8 @@ class SocketClient:
             socket_object = creator.create_socket_ipv4_tcp()
             return creator.wrap_socket_with_ssl_context_client___default_certs___ignore_verification(
                 socket_object, self.service_name, self.custom_pem_client_certificate_file_path,
-                enable_sslkeylogfile_env_to_client_ssl_context=self.enable_sslkeylogfile_env_to_client_ssl_context
+                enable_sslkeylogfile_env_to_client_ssl_context=self.enable_sslkeylogfile_env_to_client_ssl_context,
+                sslkeylog_file_path=self.sslkeylog_file_path
             )
 
     def service_connection(
@@ -215,7 +218,6 @@ class SocketClient:
         self.socket_instance = None
         self.logger.info(f"Closed socket to service server [{self.service_name}:{self.service_port}]")
 
-    # noinspection PyUnusedLocal
     def send_receive_to_service(
             self,
             request_bytes: Union[bytearray, bytes],
@@ -228,9 +230,8 @@ class SocketClient:
         :param skip_send: If True, the data will not be sent to the service server. After the connection is established,
             the function will wait for the response only.
         """
-        # Define variables
-        function_service_data = None
-        error_message = None
+
+        origin_data: bytes | None = None
 
         service_socket, error_message = self.service_connection()
         # If connection to service server wasn't successful
@@ -254,7 +255,7 @@ class SocketClient:
             if not skip_send:
                 # Send the data received from the client to the service over socket
                 error_on_send = Sender(
-                    ssl_socket=self.socket_instance, class_message=request_bytes, logger=self.logger).send()
+                    ssl_socket=self.socket_instance, bytes_to_send=request_bytes, logger=self.logger).send()
 
                 # If the socket disconnected on data send
                 if error_on_send:
@@ -265,17 +266,17 @@ class SocketClient:
 
             # Else if send was successful
             if not error_on_send:
-                function_service_data = Receiver(
+                origin_data, is_socket_closed, error_message = Receiver(
                     ssl_socket=self.socket_instance, logger=self.logger).receive()
 
                 # If data received is empty meaning the socket was closed on the other side
-                if not function_service_data:
+                if not origin_data:
                     error_message = "Service server closed the connection on receive"
 
                     # We'll close the socket and nullify the object
                     self.close_socket()
 
-        return function_service_data, error_message, self.connection_ip, self.socket_instance
+        return origin_data, error_message, self.connection_ip, self.socket_instance
 
     def send_receive_message_list_with_interval(
             self, requests_bytes_list: list, intervals_list: list, intervals_defaults: int, cycles: int = 1):

atomicshop/wrappers/socketw/socket_wrapper.py

@@ -7,17 +7,20 @@ import socket
 import shutil
 import os
 
+import paramiko
+
 from ...mitm import initialize_engines
 from ..psutilw import psutil_networks
 from ..certauthw import certauthw
 from ..loggingw import loggingw
-from ...script_as_string_processor import ScriptAsStringProcessor
+from ... import package_mains_processor
 from ...permissions import permissions
 from ... import filesystem, certificates
-from ...basics import booleans
+from ...basics import booleans, tracebacks
 from ...print_api import print_api
+from ...ssh_remote import SSHRemote
 
-from . import base, creator, get_process, accepter, statistics_csv, ssl_base, sni
+from . import socket_base, creator, process_getter, accepter, statistics_csv, ssl_base, sni
 
 
 class SocketWrapperPortInUseError(Exception):
@@ -65,10 +68,7 @@ class SocketWrapper:
             ssh_user: str = None,
             ssh_pass: str = None,
             ssh_script_to_execute: Union[
-                Literal[
-                    'process_from_port',
-                    'process_from_ipv4'
-                ],
+                Literal['process_from_port'],
                 None
             ] = None,
             logs_directory: str = None,
@@ -78,6 +78,8 @@ class SocketWrapper:
             statistics_logger_queue: multiprocessing.Queue = None,
             exceptions_logger_name: str = 'SocketWrapperExceptions',
             exceptions_logger_queue: multiprocessing.Queue = None,
+            enable_sslkeylogfile_env_to_client_ssl_context: bool = False,
+            sslkeylog_file_path: str = None,
             print_kwargs: dict = None,
     ):
         """
@@ -171,6 +173,12 @@ class SocketWrapper:
         :param exceptions_logger_name: string, name of the logger that will be used to log exceptions.
         :param exceptions_logger_queue: multiprocessing.Queue, queue that will be used to log exceptions in
             multiprocessing. You need to start the logger listener in the main process to handle the queue.
+        :param enable_sslkeylogfile_env_to_client_ssl_context: boolean, if True, each client SSL context
+            that will be created by the SocketWrapper will have save the SSL handshake keys to the file
+            defined in 'sslkeylog_file_path' parameter.
+        :param sslkeylog_file_path: string, path to file where SSL handshake keys will be saved.
+            If not provided and 'enable_sslkeylogfile_env_to_client_ssl_context' is True, then
+            the environment variable 'SSLKEYLOGFILE' will be used.
         :param print_kwargs: dict, additional arguments to pass to the print function.
         """
 
@@ -206,6 +214,9 @@ class SocketWrapper:
         self.ssh_script_to_execute = ssh_script_to_execute
         self.forwarding_dns_service_ipv4_list___only_for_localhost = (
             forwarding_dns_service_ipv4_list___only_for_localhost)
+        self.enable_sslkeylogfile_env_to_client_ssl_context: bool = (
+            enable_sslkeylogfile_env_to_client_ssl_context)
+        self.sslkeylog_file_path: str = sslkeylog_file_path
         self.print_kwargs: dict = print_kwargs
 
         self.socket_object = None
@@ -224,12 +235,17 @@ class SocketWrapper:
         # Defining listening sockets list, which will be used with "select" library in 'loop_for_incoming_sockets'.
         self.listening_sockets: list = list()
 
-        # Defining 'ssh_script_processor' variable, which will be used to process SSH scripts.
+        # Defining 'ssh_script_string' variable, which will be used to process SSH scripts.
         self.ssh_script_processor = None
         if self.get_process_name:
             # noinspection PyTypeChecker
-            self.ssh_script_processor = \
-                ScriptAsStringProcessor().read_script_to_string(self.ssh_script_to_execute)
+            self.package_processor: package_mains_processor.PackageMainsProcessor | None = package_mains_processor.PackageMainsProcessor(script_file_stem=self.ssh_script_to_execute)
+
+        else:
+            self.package_processor = None
+
+        # We will initialize it during the first 'get_process_name' function call.
+        self.ssh_client: SSHRemote | None = None
 
         # If logs directory was not set, we will use the working directory.
         if not logs_directory:
@@ -378,7 +394,8 @@ class SocketWrapper:
 
             os.makedirs(self.sni_server_certificates_cache_directory, exist_ok=True)
             print_api("Removed cached server certificates.", logger=self.logger)
-
+        else:
+            os.makedirs(self.sni_server_certificates_cache_directory, exist_ok=True)
 
         if self.install_ca_certificate_to_root_store:
             if not self.ca_certificate_filepath:
@@ -497,6 +514,13 @@ class SocketWrapper:
         listening_sockets: list = [listening_socket_object]
 
         while True:
+            engine_name: str = ''
+            source_ip: str = ''
+            source_hostname: str = ''
+            dest_port: int = 0
+            process_name: str = ''
+            domain_from_engine: str = ''
+
             try:
                 # Using "select.select" which is currently the only API function that works on all
                 # operating system types: Windows / Linux / BSD.
@@ -507,7 +531,6 @@ class SocketWrapper:
 
                 listening_ip, listening_port = listening_socket_object.getsockname()
 
-                domain_from_engine = None
                 # Get the domain to connect on this process in case on no SNI provided.
                 for domain, ip_port_dict in self.engine.domain_target_dict.items():
                     if ip_port_dict['ip'] == listening_ip:
@@ -533,41 +556,88 @@ class SocketWrapper:
 
                 self.logger.info(f"Requested domain setting: {domain_from_engine}")
 
+                engine_name = get_engine_name(domain_from_engine, [self.engine])
+
                 # Wait from any connection on "accept()".
                 # 'client_socket' is socket or ssl socket, 'client_address' is a tuple (ip_address, port).
                 client_socket, client_address, accept_error_message = accepter.accept_connection_with_error(
                     listening_socket_object, domain_from_dns_server=domain_from_engine,
                     print_kwargs={'logger': self.logger})
 
+                source_ip: str = client_address[0]
+                source_port: int = client_address[1]
+                dest_port: int = listening_socket_object.getsockname()[1]
+
+                message: str = f"Accepted connection from [{source_ip}:{source_port}] to [{listening_ip}:{dest_port}] | domain: {domain_from_engine}"
+                print_api(message, logger=self.logger)
+
+                # Not always there will be a hostname resolved by the IP address, so we will leave it empty if it fails.
+                try:
+                    source_hostname = socket.gethostbyaddr(source_ip)[0]
+                    source_hostname = source_hostname.lower()
+                except socket.herror:
+                    pass
+
                 # This is the earliest stage to ask for process name.
                 # SSH Remote / LOCALHOST script execution to identify process section.
                 # If 'get_process_name' was set to True, then this will be executed.
-                process_name = None
                 if self.get_process_name:
+                    # Initializing SSHRemote class if not initialized.
+                    if self.ssh_client is None:
+                        self.ssh_client = SSHRemote(
+                            ip_address=source_ip, username=self.ssh_user, password=self.ssh_pass, logger=self.logger)
+
                     # Get the process name from the socket.
-                    get_command_instance = get_process.GetCommandLine(
-                        client_socket=client_socket,
-                        ssh_script_processor=self.ssh_script_processor,
-                        ssh_user=self.ssh_user,
-                        ssh_pass=self.ssh_pass,
+                    get_command_instance = process_getter.GetCommandLine(
+                        client_ip=source_ip,
+                        client_port=source_port,
+                        package_processor=self.package_processor,
+                        ssh_client=self.ssh_client,
                         logger=self.logger)
                     process_name = get_command_instance.get_process_name(print_kwargs={'logger': self.logger})
 
-                source_ip: str = client_address[0]
-                engine_name: str = get_engine_name(domain_from_engine, [self.engine])
-                dest_port: int = listening_socket_object.getsockname()[1]
-
-                # Not always there will be a hostname resolved by the IP address, so we will leave it empty if it fails.
-                try:
-                    source_hostname: str = socket.gethostbyaddr(source_ip)[0]
-                except socket.herror:
-                    source_hostname = ''
+                    # from ..pywin32w.win_event_log import fetch
+                    # events = fetch.get_latest_events(
+                    #     server_ip=source_ip,
+                    #     username=self.ssh_user,
+                    #     password=self.ssh_pass,
+                    #     log_name='Security',
+                    #     count=50,
+                    #     event_id_list=[5156]
+                    # )
+                    #
+                    # source_port = client_address[1]
+                    # for event in events:
+                    #     if source_port == event['StringsDict']['Source Port']:
+                    #         process_name = event['StringsDict']['Application Name']
+                    #         break
+                    #
+                    # if process_name == '':
+                    #     raise RuntimeError("Failed to get process name from the remote host via Event Log.")
 
                 # If 'accept()' function worked well, SSL worked well, then 'client_socket' won't be empty.
                 if client_socket:
                     # Get the protocol type from the socket.
                     is_tls: bool = False
-                    tls_properties = ssl_base.is_tls(client_socket)
+
+                    try:
+                        tls_properties = ssl_base.is_tls(client_socket, timeout=1)
+                    except TimeoutError:
+                        error: str = "TimeoutError: TLS detection timed out. Dropping accepted socket."
+                        self.logger.error(error)
+
+                        self.statistics_writer.write_accept_error(
+                            engine=engine_name,
+                            source_host=source_hostname,
+                            source_ip=source_ip,
+                            error_message=error,
+                            dest_port=str(dest_port),
+                            host=domain_from_engine,
+                            process_name=process_name)
+
+                        client_socket.close()
+                        continue
+
                     if tls_properties:
                         is_tls = True
                         tls_type, tls_version = tls_properties
@@ -604,7 +674,9 @@ class SocketWrapper:
                             forwarding_dns_service_ipv4_list___only_for_localhost=(
                                 self.forwarding_dns_service_ipv4_list___only_for_localhost),
                             tls=is_tls,
-                            exceptions_logger=self.exceptions_logger
+                            exceptions_logger=self.exceptions_logger,
+                            enable_sslkeylogfile_env_to_client_ssl_context=self.enable_sslkeylogfile_env_to_client_ssl_context,
+                            sslkeylog_file_path=self.sslkeylog_file_path
                         )
 
                         ssl_client_socket, accept_error_message = \
@@ -613,6 +685,12 @@ class SocketWrapper:
                                 print_kwargs={'logger': self.logger}
                             )
 
+                        if ssl_client_socket:
+                            # Handshake is done at this point, so version/cipher are available
+                            self.logger.info(
+                                f"TLS version={ssl_client_socket.version()} cipher={ssl_client_socket.cipher()}"
+                            )
+
                         if accept_error_message:
                             # Write statistics after wrap is there was an error.
                             self.statistics_writer.write_accept_error(
@@ -660,7 +738,7 @@ class SocketWrapper:
                     self.threads_list.append(thread_current)
 
                     # 'thread_callable_args[1][0]' is the client socket.
-                    client_address = base.get_source_address_from_socket(client_socket)
+                    client_address = socket_base.get_source_address_from_socket(client_socket)
 
                     self.logger.info(f"Accepted connection, thread created {client_address}. Continue listening...")
                 # Else, if no client_socket was opened during, accept, then print the error.
@@ -674,8 +752,23 @@ class SocketWrapper:
                         dest_port=str(dest_port),
                         host=domain_from_engine,
                         process_name=process_name)
+            # Sometimes paramiko SSH connection return EOFError on connection reset, so we need to catch it separately.
+            except (ConnectionResetError, paramiko.ssh_exception.SSHException, EOFError) as e:
+                exception_string: str = tracebacks.get_as_string()
+                full_string: str = f"{str(e)} | {exception_string}"
+                self.statistics_writer.write_accept_error(
+                    engine=engine_name,
+                    source_host=source_hostname,
+                    source_ip=source_ip,
+                    error_message=full_string,
+                    dest_port=str(dest_port),
+                    host=domain_from_engine,
+                    process_name=process_name)
             except Exception as e:
-                self.exceptions_logger.write(e)
+                _ = e
+                exception_string: str = tracebacks.get_as_string()
+                full_string: str = f"Engine: [{engine_name}] | {exception_string}"
+                self.exceptions_logger.write(full_string)
 
 
 def before_socket_thread_worker(
@@ -694,7 +787,7 @@ def before_socket_thread_worker(
     try:
         callable_function(*callable_args)
     except Exception as e:
-        exceptions_logger.write(e)
+        exceptions_logger.write(e, custom_exception_attribute='engine_name', custom_exception_attribute_placement='before')
 
 
 def get_engine_name(domain: str, engine_list: list):

atomicshop/wrappers/socketw/ssl_base.py

@@ -24,15 +24,19 @@ def convert_der_x509_bytes_to_pem_string(certificate) -> str:
     return ssl.DER_cert_to_PEM_cert(certificate)
 
 
-def is_tls(client_socket) -> Optional[Tuple[str, str]]:
+def is_tls(
+        client_socket,
+        timeout: float = None
+) -> Optional[Tuple[str, str]]:
     """
     Return protocol type of the incoming socket after 'accept()'.
     :param client_socket: incoming socket after 'accept()'.
+    :param timeout: optional timeout for receiving/peeking the first bytes.
     :return: tuple with content type, protocol type + version.
         If the length of the first bytes is less than 3, return None.
     """
 
-    first_bytes = receiver.peek_first_bytes(client_socket, bytes_amount=3)
+    first_bytes = receiver.peek_first_bytes(client_socket, bytes_amount=3, timeout=timeout)
 
     # Sometimes only one byte is available, so we need to handle that case.
     # convert to a tuple of ints, add three Nones, and keep only the first 3 items.

atomicshop/wrappers/ubuntu_terminal.py

@@ -4,10 +4,15 @@ import subprocess
 import shutil
 import time
 
+from rich.console import Console
+
 from ..print_api import print_api
 from ..permissions import ubuntu_permissions
 
 
+console = Console()
+
+
 def install_packages(
         package_list: list[str],
         timeout_seconds: int = 0,
@@ -110,7 +115,7 @@ def update_system_packages():
     Function updates the system packages.
     :return:
     """
-    subprocess.check_call(['sudo', 'apt-get', 'update'])
+    subprocess.check_call(['sudo', 'apt', 'update'])
 
 
 def upgrade_system_packages(apt_update: bool = True):
@@ -124,7 +129,7 @@ def upgrade_system_packages(apt_update: bool = True):
     if apt_update:
         update_system_packages()
 
-    subprocess.check_call(['sudo', 'apt-get', 'upgrade', '-y'])
+    subprocess.check_call(['sudo', 'apt', 'upgrade', '-y'])
 
 
 def is_service_running(service_name: str, user_mode: bool = False, return_false_on_error: bool = False) -> bool:
@@ -211,28 +216,25 @@ def start_service(service_name: str, sudo: bool = False, user_mode: bool = False
 def start_enable_service_check_availability(
         service_name: str,
         wait_time_seconds: float = 30,
-        exit_on_error: bool = True,
         start_service_bool: bool = True,
         enable_service_bool: bool = True,
         check_service_running: bool = True,
         user_mode: bool = False,
-        sudo: bool = True,
-        print_kwargs: dict = None
-):
+        sudo: bool = True
+) -> int:
     """
     Function starts and enables a service and checks its availability.
 
     :param service_name: str, the service name.
     :param wait_time_seconds: float, the time to wait after starting the service before checking the service
         availability.
-    :param exit_on_error: bool, if True, the function will exit the program if the service is not available.
     :param start_service_bool: bool, if True, the service will be started.
     :param enable_service_bool: bool, if True, the service will be enabled.
     :param check_service_running: bool, if True, the function will check if the service is running.
     :param user_mode: bool, if True, the service will be started and enabled in user mode.
     :param sudo: bool, if True, the command will be executed with sudo.
-    :param print_kwargs: dict, the print arguments.
-    :return:
+
+    :return: int, 0 if the service is running, 1 if the service is not running.
     """
 
     if not start_service_bool and not enable_service_bool:
@@ -245,18 +247,19 @@ def start_enable_service_check_availability(
         enable_service(service_name, user_mode=user_mode,sudo=sudo)
 
     if check_service_running:
-        print_api(
-            f"Waiting {str(wait_time_seconds)} seconds for the program to start before availability check...",
-            **(print_kwargs or {}))
-        time.sleep(wait_time_seconds)
+        print(f"Waiting up to {str(wait_time_seconds)} seconds for the program to start.")
+        count: int = 0
+        while not is_service_running(service_name, user_mode=user_mode) and count < wait_time_seconds:
+            count += 1
+            time.sleep(1)
 
         if not is_service_running(service_name, user_mode=user_mode):
-            print_api(
-                f"[{service_name}] service failed to start.", color='red', error_type=True, **(print_kwargs or {}))
-            if exit_on_error:
-                sys.exit(1)
+            console.print(f"[{service_name}] service failed to start.", style='red', markup=False)
+            return 1
         else:
-            print_api(f"[{service_name}] service is running.", color='green', **(print_kwargs or {}))
+            console.print(f"[{service_name}] service is running.", style='green', markup=False)
+
+    return 0
 
 
 def add_path_to_bashrc(as_regular_user: bool = False):