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

atomicshop/wrappers/socketw/ssl_base.py

@@ -24,22 +24,23 @@ 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)
 
-    # Ensure we got enough data.
-    if len(first_bytes) < 3:
-        return None
-
-    # Extract the content type and version
-    content_type, version_major, version_minor = first_bytes
+    # 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.
+    content_type, version_major, version_minor = (tuple(first_bytes) + (None, None, None))[:3]
 
     # Map TLS content types to their string representation.
     content_type_map = {
@@ -52,19 +53,19 @@ def is_tls(client_socket) -> Optional[Tuple[str, str]]:
 
     # Map TLS version bytes to their string representation.
     version_map = {
-        (0x03, 0x00): "SSL 3.0",
-        (0x03, 0x01): "TLS 1.0",
-        (0x03, 0x02): "TLS 1.1",
-        (0x03, 0x03): "TLS 1.2/1.3"
+        (0x03, 0x00): "SSLv3.0",
+        (0x03, 0x01): "TLSv1.0",
+        (0x03, 0x02): "TLSv1.1",
+        (0x03, 0x03): "TLSv1.2/1.3"
         # Remember, you can't definitively differentiate 1.2 and 1.3 just from these bytes
     }
 
     # Get the tuple of the type and version as strings.
-    tls_content_and_version_tuple: tuple =\
+    tls_content_and_version_tuple: tuple[str, str] = \
         content_type_map.get(content_type), version_map.get((version_major, version_minor))
 
     # If both parts of the tuple are not None, return the protocol type.
-    if tls_content_and_version_tuple[0] and tls_content_and_version_tuple[1]:
+    if tls_content_and_version_tuple[0]:
         return tls_content_and_version_tuple
     else:
         return None

atomicshop/wrappers/socketw/statistics_csv.py

@@ -1,23 +1,154 @@
 import datetime
+import multiprocessing
 
-from ...print_api import print_api
+from ..loggingw import loggingw
 
 
-def write_accept_error(error_message: str, host: str, process_name: str, statistics_logger, print_kwargs: dict = None):
-    request_time_sent = datetime.datetime.now()
+LOGGER_NAME: str = 'statistics'
+STATISTICS_HEADER: str = (
+    'request_time_sent,thread_id,engine,source_host,source_ip,tls,protocol,protocol2,protocol3,dest_port,host,path,command,status_code,request_size_bytes,'
+    'response_size_bytes,file_path,process_cmd,action,error')
 
-    # If 'message' is not defined, it means there was no execution and there is no need for statistics.
-    try:
-        statistics_logger.info(
-            f"{request_time_sent},"
-            f"{host},"
-            f",,,,,,"
-            f"\"{process_name}\","
-            f"{error_message}"
+
+class StatisticsCSVWriter(loggingw.CsvLogger):
+    """
+    Class to write statistics to CSV file.
+    This can be initiated at the main, and then passed to the thread worker function.
+    """
+    def __init__(
+            self,
+            logger_name: str = LOGGER_NAME,
+            directory_path: str = None,
+            log_queue: multiprocessing.Queue = None,
+            add_queue_handler_start_listener_multiprocessing: bool = False,
+            add_queue_handler_no_listener_multiprocessing: bool = False
+    ):
+        """
+        Initialize the StatisticsCSVWriter with the directory path for the statistics CSV file.
+        :param directory_path: str, the directory path where the statistics CSV file will be created.
+        :param log_queue: multiprocessing.Queue, the queue to use for logging in multiprocessing.
+        :param add_queue_handler_start_listener_multiprocessing: bool, whether to add a queue handler that will use
+            the 'logger_queue' and start the queue listener with the same 'logger_queue' for multiprocessing.
+        :param add_queue_handler_no_listener_multiprocessing: bool, whether to add a queue handler that will use
+            the 'logger_queue' but will not start the queue listener for multiprocessing. This is useful when you
+            already started the queue listener and want to add more handlers to the logger without
+            starting a new listener.
+
+        If you don't set any of 'add_queue_handler_start_listener_multiprocessing' or
+        'add_queue_handler_no_listener_multiprocessing', the logger will be created without a queue handler.
+        """
+
+        super().__init__(
+            logger_name=logger_name,
+            directory_path=directory_path,
+            log_queue=log_queue,
+            add_queue_handler_start_listener_multiprocessing=add_queue_handler_start_listener_multiprocessing,
+            add_queue_handler_no_listener_multiprocessing=add_queue_handler_no_listener_multiprocessing,
+            custom_header=STATISTICS_HEADER
+        )
+
+    def write_row(
+            self,
+            thread_id: str,
+            engine: str,
+            source_host: str,
+            source_ip: str,
+            host: str,
+            tls_type: str,
+            tls_version: str,
+            protocol: str,
+            protocol2: str,
+            protocol3: str,
+            dest_port: str,
+            path: str,
+            status_code: str,
+            command: str,
+            request_size_bytes: str,
+            response_size_bytes: str,
+            recorded_file_path: str = None,
+            process_cmd: str = None,
+            error: str = None,
+            action: str = None,
+            timestamp=None,
+    ):
+        if not timestamp:
+            timestamp = datetime.datetime.now()
+
+        if not tls_type and not tls_version:
+            tls_info = ''
+        else:
+            tls_info = f'{tls_type}|{tls_version}'
+
+        row_of_cols: list = [
+            timestamp,
+            thread_id,
+            engine,
+            source_host,
+            source_ip,
+            tls_info,
+            protocol,
+            protocol2,
+            protocol3,
+            dest_port,
+            host,
+            path,
+            command,
+            status_code,
+            request_size_bytes,
+            response_size_bytes,
+            recorded_file_path,
+            process_cmd,
+            action,
+            error
+        ]
+
+        super().write(row_of_cols)
+
+    def write_accept_error(
+            self,
+            engine: str,
+            source_ip: str,
+            source_host: str,
+            error_message: str,
+            dest_port: str,
+            host: str,
+            process_name: str,
+            thread_id: str = str()
+    ):
+        """
+        Write the error message to the statistics CSV file.
+        This is used for easier execution, since most of the parameters will be empty on accept.
+
+        :param engine: string, engine name.
+        :param source_ip: string, source IP address.
+        :param source_host: string, source host name.
+        :param error_message: string, error message.
+        :param dest_port: string, destination port.
+        :param host: string, host, the domain or IP address.
+        :param process_name: process name, the command line of the process.
+        :param thread_id: integer, the id of the thread.
+        :return:
+        """
+
+        self.write_row(
+            thread_id=thread_id,
+            engine=engine,
+            source_host=source_host,
+            source_ip=source_ip,
+            tls_type='',
+            tls_version='',
+            protocol='',
+            protocol2='',
+            protocol3='',
+            dest_port=dest_port,
+            host=host,
+            path='',
+            status_code='',
+            command='',
+            request_size_bytes='',
+            response_size_bytes='',
+            recorded_file_path='',
+            process_cmd=process_name,
+            action='client_accept',
+            error=error_message
         )
-    except UnboundLocalError:
-        pass
-    except Exception:
-        message = "Undocumented exception after accept on building statistics."
-        print_api(message, error_type=True, logger_method='error', traceback_string=True, oneline=True, **print_kwargs)
-        pass

atomicshop/wrappers/sysmonw.py

@@ -0,0 +1,157 @@
+import os
+import subprocess
+
+from .. import filesystem, web, process
+
+
+# Define paths and configuration
+DEFAULT_INSTALLATION_PATH: str = 'C:\\Sysmon'
+SYSMON_FILE_NAME: str = 'Sysmon.exe'
+SYSINTERNALS_SYSMON_URL: str = 'https://download.sysinternals.com/files/Sysmon.zip'
+SYSMON_CONFIG_FILE_NAME: str = 'sysmonconfig.xml'
+SYSMON_CONFIG_FILE_PATH: str = os.path.join(DEFAULT_INSTALLATION_PATH, SYSMON_CONFIG_FILE_NAME)
+
+
+class ConfigFileNotFoundError(Exception):
+    pass
+
+
+class SymonExecutableNotFoundError(Exception):
+    pass
+
+
+class SysmonAlreadyRunningError(Exception):
+    pass
+
+
+def download_sysmon(installation_path: str = None):
+    """
+    Install Sysmon on the system.
+
+    :param installation_path: string, full path where to put the Sysmon executable.
+    """
+
+    if not installation_path:
+        installation_path = DEFAULT_INSTALLATION_PATH
+
+    # Check if the file exists
+    if not os.path.exists(installation_path):
+        filesystem.create_directory(installation_path)
+
+    web.download_and_extract_file(SYSINTERNALS_SYSMON_URL, installation_path)
+
+
+def is_sysmon_running():
+    """
+    Check if Sysmon is running.
+
+    :return: boolean, True if Sysmon is running, False otherwise.
+    """
+
+    process_list: list = process.get_running_processes_by_cmdline_pattern(
+        pattern=SYSMON_FILE_NAME, first=True, cmdline_case_insensitive=True)
+
+    if process_list:
+        return True
+    else:
+        return False
+
+
+def start_as_service(
+        installation_path: str = None,
+        config_file_path: str = None,
+        use_config_in_same_directory: bool = False,
+        download_sysmon_if_not_found: bool = False,
+        skip_if_running: bool = False
+):
+    """
+    Start Sysmon as a service. Besides starting, it installs itself as a service, meaning that on the next boot,
+    it will start automatically.
+
+    :param installation_path: string, full path where to put the Sysmon executable.
+    :param config_file_path: string, full path to the configuration file.
+    :param use_config_in_same_directory: boolean, if True, the function will use the configuration file in the same
+        directory as the Sysmon executable.
+    :param download_sysmon_if_not_found: boolean, if True, the function will download Sysmon if it is not
+        found in the 'installation_path'.
+    :param skip_if_running: boolean,
+        True, the function will not start Sysmon if it is already running.
+        False, the function will raise 'SysmonAlreadyRunningError' exception if it is already running.
+    """
+
+    # Check if sysmon already running.
+    if is_sysmon_running():
+        if skip_if_running:
+            return
+        else:
+            raise SysmonAlreadyRunningError("Sysmon is already running.")
+
+    if config_file_path and use_config_in_same_directory:
+        raise ValueError("You cannot use both 'config_file_path' and 'use_config_in_same_directory'.")
+
+    if use_config_in_same_directory:
+        config_file_path = SYSMON_CONFIG_FILE_PATH
+
+    if config_file_path:
+        # Check if the file exists
+        if not os.path.exists(config_file_path):
+            raise ConfigFileNotFoundError(f"Configuration file '{config_file_path}' not found.")
+
+    if not installation_path:
+        installation_path = DEFAULT_INSTALLATION_PATH
+
+    sysmon_file_path: str = os.path.join(installation_path, SYSMON_FILE_NAME)
+
+    # Check if the file exists
+    if not os.path.exists(sysmon_file_path):
+        if download_sysmon_if_not_found:
+            download_sysmon(installation_path)
+        else:
+            raise SymonExecutableNotFoundError(f"Sysmon executable '{sysmon_file_path}' not found.")
+
+    # Start Sysmon as a service.
+    if config_file_path:
+        subprocess.run([sysmon_file_path, '-accepteula', '-i', config_file_path])
+    else:
+        subprocess.run([sysmon_file_path, '-accepteula', '-i'])
+
+
+def stop_service(installation_path: str = None):
+    """
+    Stop Sysmon service.
+
+    :param installation_path: string, full path where to put the Sysmon executable.
+    """
+
+    if not installation_path:
+        installation_path = DEFAULT_INSTALLATION_PATH
+
+    sysmon_file_path: str = os.path.join(installation_path, SYSMON_FILE_NAME)
+
+    # Check if the file exists
+    if not os.path.exists(sysmon_file_path):
+        raise SymonExecutableNotFoundError(f"Sysmon executable '{sysmon_file_path}' not found.")
+
+    # Stop Sysmon service.
+    subprocess.run([sysmon_file_path, '-u'])
+
+
+def change_config_on_the_fly(config_file_path: str, installation_path: str = None):
+    """
+    Change the Sysmon configuration on the fly.
+
+    :param config_file_path: string, full path to the configuration file.
+    :param installation_path: string, full path where to put the Sysmon executable.
+    """
+
+    if not installation_path:
+        installation_path = DEFAULT_INSTALLATION_PATH
+
+    sysmon_file_path: str = os.path.join(installation_path, SYSMON_FILE_NAME)
+
+    # Check if the file exists
+    if not os.path.exists(sysmon_file_path):
+        raise SymonExecutableNotFoundError(f"Sysmon executable '{sysmon_file_path}' not found.")
+
+    # Change the configuration on the fly.
+    subprocess.run([sysmon_file_path, '-c', config_file_path])

atomicshop/wrappers/ubuntu_terminal.py

@@ -4,23 +4,64 @@ import subprocess
 import shutil
 import time
 
+from rich.console import Console
+
 from ..print_api import print_api
-from .. import permissions
+from ..permissions import ubuntu_permissions
+
+
+console = Console()
 
 
-def install_packages(package_list: list[str]):
+def install_packages(
+        package_list: list[str],
+        timeout_seconds: int = 0,
+):
     """
     Function installs a package using apt-get.
     :param package_list: list of strings, package names to install.
+    :param timeout_seconds: int, if the 'apt-get' command is busy at the moment, the function will wait for
+        'timeout_seconds' seconds before raising an error.
+        '-1' means wait indefinitely.
     :return:
     """
 
     # Construct the command with the package list
-    command = ["sudo", "apt-get", "install", "-y"] + package_list
+    command = ["sudo", "apt", "install", "-y"] + package_list
+
+    if timeout_seconds != 0:
+        command.extend(["-o", f"DPkg::Lock::Timeout={str(timeout_seconds)}"])
 
     subprocess.check_call(command)
 
 
+def remove_packages(
+        package_list: list[str],
+        remove_config_files: bool = False,
+        remove_dependencies: bool = False,
+):
+    """
+    Function removes a list of packages.
+    :param package_list: list of strings, package names to remove. Regular removal is through 'apt remove'.
+    :param remove_config_files: bool, if True, the config files will be removed also through 'apt purge'.
+    :param remove_dependencies: bool, if True, the dependencies will be removed also through 'apt autoremove'.
+    :return:
+    """
+
+    # Construct the command with the package list
+    command = ["sudo", "apt", "remove", "-y"] + package_list
+
+    # If remove_config_files is True, add 'purge' to the command
+    if remove_config_files:
+        command.insert(2, "purge")
+
+    subprocess.check_call(command)
+
+    # If remove_dependencies is True, remove the dependencies
+    if remove_dependencies:
+        subprocess.check_call(["sudo", "apt", "autoremove", "-y"])
+
+
 def is_package_installed(package: str) -> bool:
     """
     Function checks if a package is installed.
@@ -74,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):
@@ -88,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:
@@ -175,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:
@@ -209,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):
@@ -231,11 +270,11 @@ def add_path_to_bashrc(as_regular_user: bool = False):
 
     if as_regular_user:
         # Get the current non-sudo user
-        with permissions.temporary_regular_permissions():
+        with ubuntu_permissions.temporary_regular_permissions():
             current_non_sudo_user = os.getlogin()
 
         # Get the home path of the current non-sudo user
-        user_bashrc_path = permissions.expand_user_path(current_non_sudo_user, home_path_bashrc)
+        user_bashrc_path = ubuntu_permissions.expand_user_path(current_non_sudo_user, home_path_bashrc)
     else:
         user_bashrc_path = os.path.expanduser(home_path_bashrc)
 
@@ -258,11 +297,11 @@ def add_line_to_bashrc(line: str, as_regular_user: bool = False):
 
     if as_regular_user:
         # Get the current non-sudo user
-        with permissions.temporary_regular_permissions():
+        with ubuntu_permissions.temporary_regular_permissions():
             current_non_sudo_user = os.getlogin()
 
         # Get the home path of the current non-sudo user
-        user_bashrc_path = permissions.expand_user_path(current_non_sudo_user, home_path_bashrc)
+        user_bashrc_path = ubuntu_permissions.expand_user_path(current_non_sudo_user, home_path_bashrc)
     else:
         user_bashrc_path = os.path.expanduser(home_path_bashrc)
 
@@ -293,7 +332,7 @@ def get_command_execution_as_sudo_executer(command: str, add_bash_exec: bool = F
 
         command = command.replace('sudo ', '').strip()
 
-    sudo_executer_username: str = permissions.get_ubuntu_sudo_executer_username()
+    sudo_executer_username: str = ubuntu_permissions.get_sudo_executer_username()
 
     if sudo_executer_username:
         if add_bash_exec: