atomicshop 2.11.47__py3-none-any.whl → 3.10.5__py3-none-any.whl

atomicshop/wrappers/socketw/socket_client.py

@@ -1,6 +1,14 @@
 import socket
 import ssl
 import time
+from typing import Literal, Union
+import logging
+from pathlib import Path
+
+# noinspection PyPackageRequirements
+from cryptography import x509
+# noinspection PyPackageRequirements
+import dns.resolver
 
 from . import creator
 from .receiver import Receiver
@@ -8,19 +16,24 @@ from .sender import Sender
 from . import ssl_base
 from .. import cryptographyw
 from ..loggingw import loggingw
-from ...print_api import print_api
+from ... import print_api
 from ...file_io import file_io
-from ... import filesystem
-
-import dns.resolver
+from ...basics import tracebacks
 
 
 class SocketClient:
-    logger = loggingw.get_logger_with_level("network." + __name__.rpartition('.')[2])
-
     def __init__(
             self,
-            service_name: str, service_port: int, tls: bool = False, connection_ip=None, dns_servers_list=None):
+            service_name: str,
+            service_port: int,
+            tls: bool = False,
+            connection_ip=None,
+            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,
+            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
         domain and the connection itself (socket.connect()) made for the IP. This way YOU decide to which IPv4 your
@@ -36,6 +49,16 @@ class SocketClient:
             'service_name' by these DNS servers, with the first IPv4 result.
         :param dns_servers_list: (Optional) List object with dns IPv4 addresses that 'service_name' will be resolved
             with, using 'dnspython' module. 'connection_ip' will be populated with first resolved IP.
+        :param logger: (Optional) Logger object. If not provided, the default logger will be used.
+        :param custom_pem_client_certificate_file_path: (Optional) If specified, the SSL Socket will be created with
+            custom client certificate. The path to the file with the certificate should be provided.
+        :param enable_sslkeylogfile_env_to_client_ssl_context: boolean, enables the SSLKEYLOGFILE environment variable
+            to the SSL context. Default is False.
+            if True, SSLKEYLOGFILE will be added to SSL context with:
+            ssl_context.keylog_filename = os.environ.get('SSLKEYLOGFILE')
+            This is useful for debugging SSL/TLS connections with WireShark.
+            Since WireShark also uses this environment variable to read the key log file and apply to the SSL/TLS
+            connections, so you can see the decrypted traffic.
 
         If both 'connection_ip' and 'dns_servers_list' specified, ValueException with raise.
         """
@@ -44,6 +67,15 @@ class SocketClient:
         self.tls: bool = tls
         self.connection_ip = connection_ip
         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.
+            self.logger: logging.Logger = loggingw.get_logger_with_level(f'{logger.name}.{Path(__file__).stem}')
+        else:
+            self.logger: logging.Logger = logger
 
         self.socket_instance = None
 
@@ -62,16 +94,31 @@ class SocketClient:
     def create_service_socket(self):
         # If TLS is enabled.
         if not self.tls:
-            self.logger.info(f"Creating non-SSL socket to [{self.service_name}:{self.service_port}]")
+            log_message: str = f"Creating non-SSL socket to [{self.service_name}:{self.service_port}]"
+            print_api.print_api(log_message, logger=self.logger, logger_method='info')
             return creator.create_socket_ipv4_tcp()
         else:
-            self.logger.info(f"Creating SSL socket to [{self.service_name}:{self.service_port}]")
+            log_message: str = f"Creating SSL socket to [{self.service_name}:{self.service_port}]"
+            print_api.print_api(log_message, logger=self.logger, logger_method='info')
             socket_object = creator.create_socket_ipv4_tcp()
             return creator.wrap_socket_with_ssl_context_client___default_certs___ignore_verification(
-                socket_object, self.service_name)
+                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,
+                sslkeylog_file_path=self.sslkeylog_file_path
+            )
+
+    def service_connection(
+            self
+    ) -> tuple[
+            Union[socket.socket, ssl.SSLSocket, None],
+            Union[str, None]]:
+        """
+        Function to establish connection to server
 
-    def service_connection(self):
-        """ Function to establish connection to server """
+        :return: Tuple with socket object and error string.
+            If connection was successful, the error string will be None.
+            If connection wasn't successful, the socket object will be None.
+        """
         # Check if socket to service domain exists.
         # If not
         if not self.socket_instance:
@@ -90,8 +137,8 @@ class SocketClient:
                 f"Socket already defined to [{self.service_name}:{self.service_port}]. "
                 f"Should be connected - Reusing.")
             # Since, restart the function each send_receive iteration, and there's still a connection we need to
-            # set it True, or the socket object will be nullified in the next step.
-            return True
+            # return the socket, or the socket object will be nullified in the next step.
+            return self.socket_instance, None
 
         # If 'dns_servers_list' was provided, we will resolve the domain to ip through these servers.
         if self.dns_servers_list:
@@ -109,11 +156,20 @@ class SocketClient:
                 # Get only the first entry of the list of IPs [0]
                 self.connection_ip = function_server_address[0].to_text()
                 self.logger.info(f"Resolved to [{self.connection_ip}]")
-            except dns.resolver.NXDOMAIN:
-                self.logger.error(f"Domain {self.service_name} doesn't exist - Couldn't resolve with "
-                                  f"{self.dns_servers_list}.")
-                pass
-                return None
+            except dns.resolver.NXDOMAIN as e:
+                exception_type: str = type(e).__name__
+                error_string = (
+                    f"Socket Client Connect: {exception_type}: "
+                    f"Domain {self.service_name} doesn't exist - Couldn't resolve with {self.dns_servers_list}.")
+                print_api.print_api(error_string, logger=self.logger, logger_method='error')
+                return None, error_string
+            except dns.resolver.LifetimeTimeout as e:
+                exception_type: str = type(e).__name__
+                error_string = (
+                    f"Socket Client Connect: {exception_type}: "
+                    f"Timeout while resolving domain {self.service_name} with {self.dns_servers_list}.")
+                print_api.print_api(error_string, logger=self.logger, logger_method='error')
+                return None, error_string
 
         # If DNS was resolved correctly or DNS servers weren't specified - we can try connecting.
         # If 'connection_ip' was manually specified or resolved with 'dnspython' - the connection
@@ -128,49 +184,31 @@ class SocketClient:
         try:
             # "connect()" to the server using address and port
             self.socket_instance.connect((destination, self.service_port))
-        except ConnectionRefusedError:
-            message = f"Couldn't connect to: {self.service_name}. The server is unreachable - Connection refused."
-            print_api(message, logger=self.logger, logger_method='error', traceback_string=True, oneline=True)
-            # Socket close will be handled in the thread_worker_main
-            pass
-            return None
-        except ConnectionAbortedError:
-            message = f"Connection was aborted (by the software on host) to {self.service_name}."
-            print_api(message, logger=self.logger, logger_method='error', traceback_string=True, oneline=True)
-            # Socket close will be handled in the thread_worker_main
-            pass
-            return None
-        except socket.gaierror:
-            message = f"Couldn't resolve [{self.service_name}] to IP using default methods. " \
-                      f"Domain doesn't exist or there's no IP assigned to it."
-            print_api(message, logger=self.logger, logger_method='error', traceback_string=True, oneline=True)
-            # Socket close will be handled in the thread_worker_main
-            pass
-            return None
-        except ssl.SSLError:
-            message = f"SSLError raised on connection to {self.service_name}."
-            print_api(message, logger=self.logger, logger_method='error', traceback_string=True, oneline=True)
-            # Socket close will be handled in the thread_worker_main
-            pass
-            return None
-        except TimeoutError:
-            message = f"TimeoutError raised on connection to {self.service_name}."
-            print_api(message, logger=self.logger, logger_method='error', traceback_string=True, oneline=True)
-            # Socket close will be handled in the thread_worker_main
-            pass
-            return None
-        except ValueError as e:
-            message = f'{str(e)} | on connect to [{self.service_name}].'
-            print_api(message, logger=self.logger, logger_method='error')
-            # Socket close will be handled in the thread_worker_main
-            pass
-            return None
+        except Exception as e:
+            exception_type: str = type(e).__name__
+            exception_error: str = tracebacks.get_as_string(one_line=True)
+            error_string: str = f"Socket Client Connect: {destination}: {exception_type}"
+
+            if exception_type in ['ConnectionRefusedError', 'ConnectionAbortedError', 'ConnectionResetError',
+                                  'TimeoutError'] or 'ssl' in exception_type.lower():
+                error_message: str = f"{error_string}: {exception_error}"
+                print_api.print_api(error_message, logger=self.logger, logger_method='error')
+                return None, error_message
+            elif exception_type == 'socket.gaierror':
+                custom_error_message: str = (
+                    f"Couldn't resolve [{self.service_name}] to IP using default methods. "
+                    f"Domain doesn't exist or there's no IP assigned to it.")
+                error_message: str = f"{error_string}: {custom_error_message}"
+                print_api.print_api(error_message, logger=self.logger, logger_method='error')
+                return None, error_message
+            else:
+                raise e
 
         # If everything was fine, we'll log the connection.
         self.logger.info("Connected...")
 
         # Return the connected socket.
-        return self.socket_instance
+        return self.socket_instance, None
 
     def get_socket(self):
         return self.socket_instance
@@ -180,18 +218,25 @@ 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: bytearray):
-        # Define variables
-        function_service_data = None
-        error_string = None
+    def send_receive_to_service(
+            self,
+            request_bytes: Union[bytearray, bytes],
+            skip_send: bool = False
+    ):
+        """
+        Function to send data to service server and receive response.
 
-        # If connection to service server wasn't successful
-        if not self.service_connection():
-            error_string = "Wasn't able to connect to service, closing the destination service socket"
-            print_api(error_string, logger=self.logger, logger_method='error')
+        :param request_bytes: The data that will be sent to the service server.
+        :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.
+        """
+
+        origin_data: bytes | None = None
 
-            # We'll close the socket and nullify the object
+        service_socket, error_message = self.service_connection()
+        # If connection to service server wasn't successful
+        if error_message:
+            # Wasn't able to connect to service, closing the destination service socket and nullify the object.
             self.close_socket()
         # If the connection to the service was successful
         else:
@@ -205,27 +250,33 @@ class SocketClient:
             self.logger.info(
                 f"[{self.service_name}] resolves to ip: [{self.connection_ip}]. Pulled IP from the socket.")
 
-            # Send the data received from the client to the service over socket
-            function_data_sent = Sender(self.socket_instance, request_bytes).send()
+            # noinspection PyTypeChecker
+            error_on_send: str = None
+            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, bytes_to_send=request_bytes, logger=self.logger).send()
 
-            # If the socket disconnected on data send
-            if not function_data_sent:
-                error_string = "Service socket closed on data send"
+                # If the socket disconnected on data send
+                if error_on_send:
+                    error_message = f"Service socket closed on data send: {error_on_send}"
+
+                    # We'll close the socket and nullify the object
+                    self.close_socket()
 
-                # We'll close the socket and nullify the object
-                self.close_socket()
             # Else if send was successful
-            else:
-                function_service_data = Receiver(self.socket_instance).receive()
+            if not error_on_send:
+                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:
-                    error_string = "Service server closed the connection on receive"
+                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_string, 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):
@@ -298,15 +349,15 @@ class SocketClient:
                 # be passed.
                 # If there was connection error or socket close, then "ssl_socket" of the "service_client"
                 # will be empty.
-                response_raw_bytes, error_string, self.connection_ip, service_ssl_socket = \
+                response_raw_bytes, error_message, self.connection_ip, service_ssl_socket = \
                     self.send_receive_to_service(request_raw_bytes)
 
                 # Adding the response to responses list. Same for error.
                 responses_list.append(response_raw_bytes)
-                errors_list.append(error_string)
+                errors_list.append(error_message)
 
                 self.logger.info(f"Response: {response_raw_bytes}")
-                self.logger.info(f"Error: {error_string}")
+                self.logger.info(f"Error: {error_message}")
 
                 # So if the socket was closed and there was an error we can break the loop.
                 # This is needed for more complex operations
@@ -323,9 +374,9 @@ class SocketClient:
             self,
             save_as_file: bool = False,
             cert_file_path: str = None,
-            cert_output_type: str = 'der',
+            cert_output_type: Literal['der', 'cryptography'] = 'der',
             **kwargs
-    ):
+    ) -> Union[x509.Certificate]:
         """
         This function will get the certificate from the server and return it.
 
@@ -339,21 +390,12 @@ class SocketClient:
         # If "save_as_file" is True, then "cert_file_path" must be provided, if not, raise an exception.
         if save_as_file and not cert_file_path:
             raise ValueError("If 'save_as_file' is True, then 'cert_file_path' must be provided.")
-        # If 'save_as_file' is True and 'cert_file_path' is provided, then check if the file exists.
-        # Since there is no point fetching the certificate from the socket if it already exists.
-        elif save_as_file and cert_file_path:
-            # If certificate from socket exists, then we don't need to get it from the socket and write to file.
-            # and we will return None, since no certificate was fetched.
-            if filesystem.check_file_existence(cert_file_path):
-                return None
-            else:
-                print_api("Certificate from socket doesn't exist, fetching.", logger=self.logger)
 
         # Connect and get the connected socket.
-        server_socket_for_certificate = self.service_connection()
+        server_socket_for_certificate, error_message = self.service_connection()
         # Get the DER byte certificate from the socket.
         certificate_from_socket_der_bytes = ssl_base.get_certificate_from_socket(server_socket_for_certificate)
-        print_api('Fetched certificate from socket.', logger=self.logger, **kwargs)
+        print_api.print_api('Fetched certificate from socket.', logger=self.logger, **kwargs)
         # Close the socket.
         self.close_socket()
 

atomicshop/wrappers/socketw/socket_server_tester.py

@@ -3,7 +3,7 @@ import threading
 from .socket_client import SocketClient
 from ..configparserw import ConfigParserWrapper
 from ..loggingw import loggingw
-from ...filesystem import get_file_paths_from_directory
+from ... import filesystem
 from ...file_io import jsons, file_io
 
 
@@ -37,7 +37,7 @@ def get_key_values_from_json(json_dict: dict, extract_keys: list):
 def execute_test(config_static):
     # Import config ini file and read it to dict.
     config_importer = ConfigParserWrapper(
-        file_name=config_static.CONFIG_INI_TESTER_FILE_NAME, directory_path=config_static.WORKING_DIRECTORY)
+        file_name=config_static.CONFIG_INI_TESTER_FILE_NAME, directory_path=config_static.SCRIPT_DIRECTORY)
     config_importer.read_to_dict()
     # Convert keys.
     config_importer.convert_string_values(
@@ -54,17 +54,21 @@ def execute_test(config_static):
     config = config_importer.config['config']
 
     # SocketClient is working with 'network' logger by default, so we will initialize it.
-    loggingw.get_logger_with_stream_handler("network")
+    loggingw.create_logger(
+        logger_name="network",
+        add_stream=True,
+        formatter_streamhandler='DEFAULT'
+    )
 
     # Get all the files in requests folder recursively.
-    request_file_list = get_file_paths_from_directory(config['requests_directory'])
+    request_file_list = filesystem.get_paths_from_directory(config['requests_directory'], get_file=True)
     print(f"Found request files: {len(request_file_list)}")
 
     # Get contents of all request files to list of contents.
     requests_bytes_list: list = list()
     for request_file_path in request_file_list:
         if config['request_type'] == 'json':
-            request_file_content = jsons.read_json_file(request_file_path)
+            request_file_content = jsons.read_json_file(request_file_path.path)
 
             # If imported json is regular and not combined json.
             if isinstance(request_file_content, dict):
@@ -79,13 +83,13 @@ def execute_test(config_static):
                     requests_bytes_list.extend(
                         get_key_values_from_json(json_dict, config['request_json_hex_key_list']))
         elif config['request_type'] == 'string':
-            request_file_content = file_io.read_file(request_file_path)
+            request_file_content = file_io.read_file(request_file_path.path)
             # Convert string content to bytes and append to list.
             requests_bytes_list.append(request_file_content.encode())
             print(f"Extracted 1 request.")
         elif config['request_type'] == 'binary':
             # The content is already in bytes, so just appending.
-            requests_bytes_list.append(file_io.read_file(request_file_path, 'rb'))
+            requests_bytes_list.append(file_io.read_file(request_file_path.path, 'rb'))
             print(f"Extracted 1 request.")
 
     print(f"Finished parsing. Parsed requests: {len(requests_bytes_list)}")
@@ -117,10 +121,11 @@ def execute_test(config_static):
                     requests_bytes_list,
                     config['interval_between_request_custom_list_seconds'],
                     config['interval_between_requests_defaults_seconds'],))
-            # Append to list of threads, so they can be "joined" later
-            threads_list.append(thread_current)
+            thread_current.daemon = True
             # Start the thread
             thread_current.start()
+            # Append to list of threads, so they can be "joined" later
+            threads_list.append(thread_current)
 
         # Joining all the threads.
         for thread in threads_list: