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

atomicshop/print_api.py

@@ -1,17 +1,18 @@
 import sys
+import logging
+from typing import Any
 
-from .basics.ansi_escape_codes import ColorsBasic, get_colors_basic_dict
+from .basics import ansi_escape_codes
 from .basics import tracebacks
 
 
-# noinspection PyUnusedLocal,PyIncorrectDocstring
 def print_api(
-        message: any,
-        color: any = None,
+        message: Any,
+        color: Any = None,
         print_end: str = '\n',
         rtl: bool = False,
         error_type: bool = False,
-        logger: object = None,
+        logger: logging.Logger = None,
         logger_method: str = 'info',
         stdout: bool = True,
         stderr: bool = True,
@@ -19,8 +20,6 @@ def print_api(
         traceback_string: bool = False,
         oneline: bool = False,
         oneline_end: str = '',
-        # raise_exception: bool = True,
-        # exit_on_error: bool = False,
         **kwargs: object) -> None:
     """
     Function of custom api that is responsible for printing messages to console.
@@ -82,39 +81,37 @@ def print_api(
 
         # This section takes care of different types of string manipulations for message.
 
-        # If 'exit_on_error' is set to 'True', we'll add 'exit_message' on new line after 'message'.
-        # if error_type and exit_on_error and raise_exception:
-        #     message = message + '\n' + exit_message
-
         # If 'rtl' is set to 'True', we'll add Right-To-Left text conversion to 'message'.
         if rtl:
-            # Lazy importing of 'bidi' library. It's not a problem since python caches the library after first import.
-            # Off-course, it will be imported from the cache each time this section is triggered.
             # pip install python-bidi
             from bidi.algorithm import get_display
             message = get_display(message)
 
+        if logger_method == 'error' or logger_method == 'critical':
+            error_type = True
+
+        # If exception was raised and 'stderr=True'.
+        if sys.exc_info()[0] is not None and stderr and traceback_string:
+            # If 'traceback' is set to 'True', we'll output traceback of exception.
+            if traceback_string:
+                if message:
+                    message = f'{message}\n{tracebacks.get_as_string()}{message}'
+                else:
+                    message = tracebacks.get_as_string()
+
+            color = 'red'
+
         # If 'stdcolor' is 'True', the console output will be colored.
         if stdcolor:
             # If 'logger.error' should be outputted to console, and 'color' wasn't selected, then set color to 'yellow'.
             if logger_method == 'error' and not color:
                 color = 'yellow'
-                error_type = True
             # If 'logger.critical' should be outputted to console, and 'color' wasn't selected, then set color to 'red'.
             elif logger_method == 'critical' and not color:
                 color = 'red'
-                error_type = True
-
-            if color:
-                message = get_colors_basic_dict(color) + message + ColorsBasic.END
 
-        # If exception was raised and 'stderr=True'.
-        if sys.exc_info()[0] is not None and stderr:
-            # If 'traceback' is set to 'True', we'll output traceback of exception.
-            if traceback_string:
-                message = f'{message} | Exception: {tracebacks.get_as_string()}'
-
-            color = 'red'
+            if color is not None:
+                message = ansi_escape_codes.get_colors_basic_dict(color) + message + ansi_escape_codes.ColorsBasic.END
 
         # If 'online' is set to 'True', we'll output message as oneline.
         if oneline:
@@ -129,14 +126,14 @@ def print_api(
             if print_end == '\n':
                 # Use logger to output message.
                 getattr(logger, logger_method)(message)
+            else:
+                raise ValueError("Logger can't output messages with 'print_end' other than '\\n'.")
         # If logger wasn't passed.
         else:
             # Use print to output the message.
             print(message, end=print_end)
 
     # = Main Section with printing cases ===============================================================================
-    # exit_message: str = 'Exiting...'
-
     # Convert message to string.
     message = str(message)
 
@@ -150,20 +147,6 @@ def print_api(
         if error_type:
             print_or_logger()
 
-    # ==================================
-    # This section is responsible for ending the script.
-
-    # Check if we're inside exception. In this case each of 3 entries in 'sys.exc_info()' tuple will not equal
-    # to 'None', so picked only the first one.
-    # if sys.exc_info()[0] and not exit_on_error:
-    #     # If 'raise_exception' is set to 'True', we'll end the script with exception.
-    #     if pass_exception:
-    #         pass
-
-    # If 'exit_on_error' is set to 'True', we'll end the script.
-    # if exit_on_error and error_type:
-    #     sys.exit()
-
 
 def print_status(
         prefix_string: str,

atomicshop/process.py

@@ -7,11 +7,16 @@ import shutil
 
 from .print_api import print_api
 from .inspect_wrapper import get_target_function_default_args_and_combine_with_current
-from .basics.strings import match_pattern_against_string
+from .basics import strings
 from .wrappers import ubuntu_terminal
+from .wrappers.psutilw import processes
 
 if os.name == 'nt':
-    from .process_poller import GetProcessList
+    from . import get_process_list
+
+
+class MultipleProcessesFound(Exception):
+    pass
 
 
 def is_command_exists(cmd: str) -> bool:
@@ -77,7 +82,18 @@ def execute_with_live_output(
     :return: Boolean, If execution was successful, return True, if not - False.
     """
 
-    cmd = _execution_parameters_processing(cmd, wsl)
+    if isinstance(cmd, str):
+        shell = True
+    elif isinstance(cmd, list):
+        shell = False
+    else:
+        raise TypeError(f'cmd must be a string or list, not {type(cmd)}')
+
+    if wsl:
+        if isinstance(cmd, str):
+            cmd = 'wsl ' + cmd
+        elif isinstance(cmd, list):
+            cmd = ['wsl'] + cmd
 
     # Needed imports:
     # from subprocess import Popen, PIPE, STDOUT
@@ -98,7 +114,7 @@ def execute_with_live_output(
     #   The buffer size is system-dependent and usually chosen by the underlying implementation to optimize performance.
     #   # bufsize=0: This means no buffering.
     #   The I/O is unbuffered, and data is written or read from the stream immediately.
-    with subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, bufsize=1, text=True) as process:
+    with subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, bufsize=1, text=True, shell=shell) as process:
         # We'll count the number of lines from 'process.stdout'.
         counter: int = 0
         # And also get list of all the lines.
@@ -247,44 +263,74 @@ def safe_terminate(popen_process: subprocess.Popen):
     popen_process.wait()
 
 
-def match_pattern_against_running_processes_cmdlines(pattern: str, first: bool = False, prefix_suffix: bool = False):
+def get_running_processes_by_cmdline_pattern(
+        pattern: str,
+        cmdline_case_insensitive: bool = False,
+        first: bool = False,
+        prefix_suffix: bool = False
+) -> list:
     """
     The function matches specified string pattern including wildcards against all the currently running processes'
     command lines.
 
     :param pattern: string, the pattern that we will search in the command line list of currently running processes.
+    :param cmdline_case_insensitive: boolean,
+        True, the pattern and the command line will be matched case-insensitive.
     :param first: boolean, that will set if first pattern match found the iteration will stop, or we will return
         the list of all command lines that contain the pattern.
     :param prefix_suffix: boolean. Check the description in 'match_pattern_against_string' function.
+
+    :return: list, of command lines that contain the pattern.
     """
 
     # Get the list of all the currently running processes.
-    get_process_list = GetProcessList(get_method='pywin32', connect_on_init=True)
-    processes = get_process_list.get_processes(as_dict=False)
+    get_process_list_instance = get_process_list.GetProcessList(get_method='psutil', connect_on_init=True)
+    processes = get_process_list_instance.get_processes(as_dict=False)
 
     # Iterate through all the current process, while fetching executable file 'name' and the command line.
     # Name is always populated, while command line is not.
     matched_cmdlines: list = list()
     for process in processes:
         # Check if command line isn't empty and that string pattern is matched against command line.
-        if process['cmdline'] and \
-                match_pattern_against_string(pattern, process['cmdline'], prefix_suffix):
-            matched_cmdlines.append(process['cmdline'])
-            # If 'first' was set to 'True' we will stop, since we found the first match.
-            if first:
-                break
+        if process['cmdline']:
+            is_pattern_matched = strings.match_pattern_against_string(
+                pattern, process['cmdline'], case_insensitive=cmdline_case_insensitive, prefix_suffix=prefix_suffix)
+            if is_pattern_matched:
+                matched_cmdlines.append(process)
+                # If 'first' was set to 'True' we will stop, since we found the first match.
+                if first:
+                    break
 
     return matched_cmdlines
 
 
-def run_powershell_command(command):
+def kill_process_by_filename_pattern(pattern: str):
+    running_processes = get_running_processes_by_cmdline_pattern(
+        pattern, first=False, prefix_suffix=True, cmdline_case_insensitive=True)
+
+    processes_found: int = len(running_processes)
+
+    if processes_found > 1:
+        raise MultipleProcessesFound(f"[{processes_found}] processes found with pattern '{pattern}'.")
+    elif processes_found == 1:
+        processes.kill_process_by_pid(running_processes[0]['pid'])
+
+
+def run_powershell_command(
+        command: str
+):
     try:
         result = subprocess.run(["powershell", "-Command", command], capture_output=True, text=True, check=True)
-        print_api(result.stdout)
+        if result.stdout:
+            print_api(result.stdout)
+        if result.stderr:  # PS can write warnings to stderr even on success
+            print_api(result.stderr, color='yellow')
         return result.stdout
     except subprocess.CalledProcessError as e:
-        print_api(f"An error occurred: {e}", color='red', error_type=True)
-        return e
+        # e.stderr and e.stdout are populated because capture_output=True
+        msg = (e.stderr or e.stdout or f"PowerShell exited with code {e.returncode}")
+        print_api(msg, color='red', error_type=True)
+        return msg
 
 
 """

atomicshop/process_poller/pollers/psutil_pywin32wmi_dll.py

@@ -0,0 +1,95 @@
+from typing import Union, Literal
+import threading
+import time
+
+from ... import get_process_list
+
+
+class PollerPsutilPywin32Dll:
+    """
+    The class is responsible for getting the list of opened processes by using mentioned libraries.
+    """
+    def __init__(
+            self,
+            interval_seconds: Union[int, float] = 0,
+            process_get_method: Literal[
+                'poll_psutil',
+                'poll_pywin32',
+                'poll_process_dll'
+            ] = 'poll_process_dll',
+            process_queue=None
+    ):
+        """
+        :param interval_seconds: works only for pollers, float, how many seconds to wait between each cycle.
+            Default is 0, which means that the polling will be as fast as possible.
+
+            Basically, you want it to be '0' if you want to get the most recent processes.
+            Any polling by itself takes time, so if you want to get the most recent processes, you want to do it as fast
+            as possible.
+        :param process_get_method: str. Default is 'process_dll'. Available:
+            'poll_psutil': Poller, Get the list of processes by 'psutil' library. Resource intensive and slow.
+            'poll_pywin32': Poller, processes by 'pywin32' library, using WMI. Not resource intensive, but slow.
+            'poll_process_dll'. Poller, Not resource intensive and fast. Probably works only in Windows 10 x64.
+            'trace_sysmon_etw': Tracer, Get the list of processes with running SysMon by ETW - Event Tracing.
+                In this case 'interval_seconds' is irrelevant, since the ETW is real-time.
+                Steps we take:
+                    1. Check if SysMon is Running. If not, check if the executable exists in specified
+                        location and start it as a service.
+                    2. Start the "Microsoft-Windows-Sysmon" ETW session.
+                    3. Take a snapshot of current processes and their CMDs with psutil and store it in a dict.
+                    4. Each new process creation from ETW updates the dict.
+            'trace_event_log': Get the list of processes by subscribing to the Windows Event Log.
+                Log Channel: Security, Event ID: 4688.
+                We enable the necessary prerequisites in registry and subscribe to the event.
+        :param process_queue: Queue. The queue to put the processes in. If None, the processes will not be put in the
+            queue.
+        """
+
+        self.interval_seconds: Union[int, float] = interval_seconds
+        self.process_get_method = process_get_method.replace('poll_', '')
+        self.process_queue = process_queue
+
+        # noinspection PyTypeChecker
+        self.poller_instance = get_process_list.GetProcessList(get_method=self.process_get_method)
+
+        self._processes = {}
+
+    def start(self):
+        """
+        Start the poller.
+        """
+
+        thread = threading.Thread(target=self.emit_loop)
+        thread.daemon = True
+        thread.start()
+
+    def emit_loop(self):
+        """
+        Get the list of processes.
+        """
+
+        self.poller_instance.connect()
+
+        while True:
+            current_processes = self.poller_instance.get_processes(as_dict=True)
+
+            # Remove Command lines that contains only numbers, since they are useless.
+            for pid, process_info in current_processes.items():
+                if process_info['cmdline'].isnumeric():
+                    current_processes[pid]['cmdline'] = str()
+                elif process_info['cmdline'] == 'Error':
+                    current_processes[pid]['cmdline'] = str()
+
+            # This loop is essential for keeping the command lines.
+            # When the process unloads from memory, the last polling will have only pid and executable name, but not
+            # the command line. This loop will keep the command line from the previous polling if this happens.
+            for pid, process_info in current_processes.items():
+                if pid in self._processes:
+                    if self._processes[pid]['name'] == current_processes[pid]['name']:
+                        if current_processes[pid]['cmdline'] == '':
+                            current_processes[pid]['cmdline'] = self._processes[pid]['cmdline']
+            self._processes.update(current_processes)
+
+            self.process_queue.put(self._processes)
+
+            time.sleep(self.interval_seconds)

atomicshop/process_poller/process_pool.py

@@ -0,0 +1,207 @@
+import threading
+import multiprocessing
+import time
+from typing import Literal, Union
+
+from .tracers import sysmon_etw, event_log
+from .pollers import psutil_pywin32wmi_dll
+from ..wrappers.pywin32w.win_event_log.subscribes import process_terminate
+
+
+POLLER_SETTINGS_SYSMON_ETW: dict = {
+    'etw_session_name': None,           # Default will be used, check 'trace_sysmon_process_creation' for details.
+    'sysmon_directory': None            # Directory, where sysmon.exe is located. Default will be used.
+}
+
+
+class ProcessPool:
+    """
+    The class is responsible for getting processes and storing them in a pool.
+    THE POOL OF PROCESSES IS NOT REAL TIME!!!
+    There can be several moments delay (less than a second) + you can add delay before pid removal from the pool.
+    This is needed only to correlate PIDs to process names and command lines of other events you get on Windows.
+    """
+    def __init__(
+            self,
+            operation: Literal['thread', 'process'] = 'thread',
+            interval_seconds: Union[int, float] = 0,
+            process_get_method: Literal[
+                'poll_psutil',
+                'poll_pywin32',
+                'poll_process_dll',
+                'trace_sysmon_etw',
+                'trace_event_log'
+            ] = 'trace_event_log',
+            poller_settings: dict = None,
+            process_terminator: bool = True,
+            wait_before_pid_remove_seconds: float = 5
+    ):
+        """
+        :param operation: str, 'thread' or 'process'. Default is 'process'.
+            'process': The polling will be done in a new process.
+            'thread': The polling will be done in a new thread.
+
+            Python is slow, if you are going to use 'thread' all other operations inside this thread will be very slow.
+            You can even get exceptions, if you're using process polling for correlations of PIDs and process names.
+            It is advised to use the 'process' operation, which will not affect other operations in the thread.
+        :param interval_seconds: works only for pollers, float, how many seconds to wait between each cycle.
+            Default is 0, which means that the polling will be as fast as possible.
+
+            Basically, you want it to be '0' if you want to get the most recent processes.
+            Any polling by itself takes time, so if you want to get the most recent processes, you want to do it as fast
+            as possible.
+        :param process_get_method: str. Default is 'process_dll'. Available:
+            'poll_psutil': Poller, Get the list of processes by 'psutil' library. Resource intensive and slow.
+            'poll_pywin32': Poller, processes by 'pywin32' library, using WMI. Not resource intensive, but slow.
+            'poll_process_dll'. Poller, Not resource intensive and fast. Probably works only in Windows 10 x64.
+            'trace_sysmon_etw': Tracer, Get the list of processes with running SysMon by ETW - Event Tracing.
+                In this case 'interval_seconds' is irrelevant, since the ETW is real-time.
+                Steps we take:
+                    1. Check if SysMon is Running. If not, check if the executable exists in specified
+                        location and start it as a service.
+                    2. Start the "Microsoft-Windows-Sysmon" ETW session.
+                    3. Take a snapshot of current processes and their CMDs with psutil and store it in a dict.
+                    4. Each new process creation from ETW updates the dict.
+            'trace_event_log': Get the list of processes by subscribing to the Windows Event Log.
+                Log Channel: Security, Event ID: 4688.
+                We enable the necessary prerequisites in registry and subscribe to the event.
+        :param poller_settings: dict, settings for the poller method.
+            'sysmon_etw': If not set 'POLLER_SETTINGS_SYSMON_ETW' will be used.
+        :param process_terminator: bool, if True, process terminator will run in the background and monitor
+            the processes for termination. If the process is terminated it will be removed from the pool.
+        :param wait_before_pid_remove_seconds: float, how many seconds to wait before the process is removed from
+            the pool after process termination event is received for that pid.
+        ---------------------------------------------
+        If there is an exception, ProcessPollerPool.processes will be set to the exception.
+        While getting the processes you can use this to execute the exception:
+
+        processes = ProcessPollerPool.processes
+
+        if isinstance(processes, BaseException):
+            raise processes
+        """
+
+        self.operation: str = operation
+        self.interval_seconds: float = interval_seconds
+        self.process_get_method = process_get_method
+        self.process_terminator: bool = process_terminator
+        self.wait_before_pid_remove_seconds: float = wait_before_pid_remove_seconds
+        self.poller_settings: dict = poller_settings
+
+        if self.poller_settings is None:
+            if process_get_method == 'sysmon_etw':
+                self.poller_settings: dict = POLLER_SETTINGS_SYSMON_ETW
+
+        # Current process pool.
+        self._processes: dict = dict()
+
+        # The variable is responsible to stop the thread if it is running.
+        self._running: bool = False
+
+        self._process_queue = multiprocessing.Queue()
+        self._running_state_queue = multiprocessing.Queue()
+
+    def start(self):
+        if self.operation == 'thread':
+            self._start_thread()
+        elif self.operation == 'process':
+            self._start_process()
+        else:
+            raise ValueError(f'Invalid operation type [{self.operation}]')
+
+        thread_get_queue = threading.Thread(target=self._thread_get_queue)
+        thread_get_queue.daemon = True
+        thread_get_queue.start()
+
+        if self.process_terminator:
+            thread_process_termination = threading.Thread(target=self._thread_process_termination)
+            thread_process_termination.daemon = True
+            thread_process_termination.start()
+
+    def stop(self):
+        self._running = False
+        self._running_state_queue.put(False)
+
+    def get_processes(self):
+        return self._processes
+
+    def _get_args_for_worker(self):
+        return self.process_get_method, self.interval_seconds, self._process_queue, self.poller_settings
+
+    def _start_thread(self):
+        self._running = True
+
+        thread = threading.Thread(target=_worker, args=(self._get_args_for_worker()))
+        thread.daemon = True
+        thread.start()
+
+    def _start_process(self):
+        self._running = True
+        multiprocessing.Process(target=_worker, args=(self._get_args_for_worker())).start()
+
+    def _thread_get_queue(self):
+        while True:
+            self._processes = self._process_queue.get()
+
+    def _thread_process_termination(self):
+        process_terminate_instance = process_terminate.ProcessTerminateSubscriber()
+        process_terminate_instance.start()
+
+        while True:
+            termination_event = process_terminate_instance.emit()
+            process_id = termination_event['ProcessIdInt']
+
+            removal_thread = threading.Thread(target=self._remove_pid, args=(process_id,))
+            removal_thread.daemon = True
+            removal_thread.start()
+
+    def _remove_pid(self, process_id):
+        # We need to wait a bit before we remove the process.
+        # This is because termination event can come sooner than the creation and the process
+        # is not yet in the pool.
+        # This happens mostly when the process is terminated immediately after the creation.
+        # Example: ping example.c
+        # 'example.c' is not a valid address, so the process is terminated immediately after the creation.
+        counter = 0
+        while counter < 30:
+            if process_id in self._processes:
+                break
+            counter += 1
+            time.sleep(0.1)
+
+        if counter == 30:
+            # print_api(f'Process [{process_id}] not found in the pool.', color='yellow')
+            return
+
+        # time.sleep(1)
+        # if process_id not in self._processes:
+        #     print_api(f'Process [{process_id}] not found in the pool.', color='red')
+        #     return
+
+        time.sleep(self.wait_before_pid_remove_seconds)
+        _ = self._processes.pop(process_id, None)
+        # print_api(f'Process [{process_id}] removed from the pool.', color='yellow')
+
+
+def _worker(
+        poller_method,
+        interval_seconds,
+        process_queue,
+        poller_settings
+):
+
+    if poller_method == 'trace_sysmon_etw':
+        get_instance = sysmon_etw.TracerSysmonEtw(settings=poller_settings, process_queue=process_queue)
+    elif poller_method == 'trace_event_log':
+        get_instance = event_log.TracerEventlog(process_queue=process_queue)
+    elif 'poll_' in poller_method:
+        get_instance = psutil_pywin32wmi_dll.PollerPsutilPywin32Dll(
+            interval_seconds=interval_seconds, process_get_method=poller_method, process_queue=process_queue)
+    else:
+        raise ValueError(f'Invalid poller method [{poller_method}]')
+
+    # We must initiate the connection inside the thread/process, because it is not thread-safe.
+    get_instance.start()
+
+    while True:
+        time.sleep(1)