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

atomicshop/process_poller/simple_process_pool.py

@@ -0,0 +1,311 @@
+import threading
+from pathlib import Path
+import time
+import multiprocessing.managers
+import queue
+
+
+from ..wrappers.pywin32w.win_event_log.subscribes import process_create, process_terminate
+from .. import get_process_list
+from ..print_api import print_api
+
+
+WAIT_BEFORE_PROCESS_TERMINATION_CHECK_SECONDS: float = 3
+WAIT_BEFORE_PROCESS_TERMINATION_CHECK_COUNTS: float = WAIT_BEFORE_PROCESS_TERMINATION_CHECK_SECONDS * 10
+
+WAIT_FOR_PROCESS_POLLER_PID_SECONDS: int = 3
+WAIT_FOR_PROCESS_POLLER_PID_COUNTS: int = WAIT_FOR_PROCESS_POLLER_PID_SECONDS * 10
+
+
+class PidProcessConverterPIDNotFoundError(Exception):
+    pass
+
+
+class SimpleProcessPool:
+    """
+    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.
+    The idea is similar to the process_poller.process_pool.ProcessPool class, but this class is simpler and uses
+    only the pywin32 tracing of the Windows Event Log Process Creation and Process Termination events.
+    The simple process pool is used to get things simpler than the process_pool.ProcessPool class.
+
+    Example of starting the process pool in multiprocess:
+        import sys
+
+        from atomicshop.process_poller import simple_process_pool
+
+
+        def start_process_pool(process_pool_shared_dict_proxy):
+            process_poller = simple_process_pool.SimpleProcessPool(
+                process_pool_shared_dict_proxy=process_pool_shared_dict_proxy)
+            process_poller.start()
+
+            try:
+                # Keep the process alive.
+                while True:
+                    time.sleep(1)
+            except KeyboardInterrupt:
+                process_poller.stop()
+
+
+        def main():
+            # Create the shared multiprocessing dictionary of the process pool.
+            manager = multiprocessing.Manager()
+            multiprocess_dict_proxy = manager.dict()
+
+            # Start the process pool in a separate process.
+            pool_process = multiprocessing.Process(target=start_process_pool, args=(multiprocess_dict_proxy,))
+            pool_process.start()
+
+            # Pass the shared dict proxy to other functions.
+
+
+        if __name__ == '__main__':
+            sys.exit(main())
+    """
+
+    def __init__(
+            self,
+            wait_before_pid_remove_seconds: float = 5,
+            process_pool_shared_dict_proxy: multiprocessing.managers.DictProxy = None
+    ):
+        """
+        :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.
+        :param process_pool_shared_dict_proxy: multiprocessing.managers.DictProxy, shared dict proxy to update
+            the process pool.
+            If you run a function from other multiprocessing.Process, you can pass the shared_dict_proxy to the function
+            and update the process pool from that function.
+
+            Example:
+            import multiprocessing.managers
+
+            manager: multiprocessing.managers.SyncManager = multiprocessing.Manager()
+            multiprocess_dict_proxy: multiprocessing.managers.DictProxy = manager.dict()
+
+            process_poller = SimpleProcessPool()
+            process_poller.start()
+
+            while True:
+                #============================
+                # your function where you get info with pid
+                # result = {
+                #    'pid': 1234,
+                #    'info': 'some info'
+                # }
+                #============================
+                info_with_process_details = {
+                    'pid': result['pid'],
+                    'info': result['info']
+                    'process_details': shared_dict_proxy[result['pid']]
+                }
+
+                break
+
+            process_poller.stop()
+            manager.shutdown()
+        """
+
+        self.wait_before_pid_remove_seconds: float = wait_before_pid_remove_seconds
+        self.process_pool_shared_dict_proxy: multiprocessing.managers.DictProxy = process_pool_shared_dict_proxy
+
+        self._processes: dict = dict()
+        self._running: bool = False
+        self.shared_dict_update_queue: queue.Queue = queue.Queue()
+        self.empty_cmdline_queue: queue.Queue = queue.Queue()
+
+    def start(self):
+        self._running = True
+
+        self._processes = get_process_list.GetProcessList(
+            get_method='psutil', connect_on_init=True).get_processes(as_dict=True)
+
+        thread_get_queue = threading.Thread(target=self._start_main_thread, args=(self.shared_dict_update_queue,))
+        thread_get_queue.daemon = True
+        thread_get_queue.start()
+
+        thread_process_termination = threading.Thread(target=self._thread_process_termination)
+        thread_process_termination.daemon = True
+        thread_process_termination.start()
+
+        thread_get_psutil_commandline = threading.Thread(target=self._thread_get_psutil_commandline)
+        thread_get_psutil_commandline.daemon = True
+        thread_get_psutil_commandline.start()
+
+        thread_update_shared_dict = threading.Thread(target=self._update_shared_dict, args=(self.shared_dict_update_queue,))
+        thread_update_shared_dict.daemon = True
+        thread_update_shared_dict.start()
+
+    def stop(self):
+        self._running = False
+
+    def get_processes(self):
+        return self._processes
+
+    def _start_main_thread(self, shared_dict_update_queue):
+        get_instance = process_create.ProcessCreateSubscriber()
+        get_instance.start()
+
+        while self._running:
+            event = get_instance.emit()
+            process_id = event['NewProcessIdInt']
+            process_name = Path(event['NewProcessName']).name
+            command_line = event['CommandLine']
+
+            # The event log tracing method doesn't always give the command line, unlike the psutil method.
+            # So, we'll get the command line from the current psutil snapshot separately.
+            if command_line == '':
+                self.empty_cmdline_queue.put(process_id)
+
+            self._processes[process_id] = {
+                'name': process_name,
+                'cmdline': command_line
+            }
+
+            # Update the multiprocessing shared dict proxy.
+            shared_dict_update_queue.put(dict(self._processes))
+
+    def _thread_get_psutil_commandline(self):
+        """
+        This function will get an entry from the queue where command line is missing and get the command line
+        from the psutil snapshot.
+        """
+
+        while self._running:
+            empty_cmd_pid = self.empty_cmdline_queue.get()
+            current_psutil_process_snapshot: dict = get_process_list.GetProcessList(
+                get_method='psutil', connect_on_init=True).get_processes(as_dict=True)
+            command_line = current_psutil_process_snapshot[empty_cmd_pid]['cmdline']
+
+            self._processes[empty_cmd_pid]['cmdline'] = command_line
+
+            # Update the multiprocessing shared dict proxy.
+            self.shared_dict_update_queue.put(dict(self._processes))
+
+    def _thread_process_termination(self):
+        process_terminate_instance = process_terminate.ProcessTerminateSubscriber()
+        process_terminate_instance.start()
+
+        while self._running:
+            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 < WAIT_BEFORE_PROCESS_TERMINATION_CHECK_COUNTS:
+            if process_id in self._processes:
+                break
+            counter += 1
+            time.sleep(0.1)
+
+        if counter == WAIT_BEFORE_PROCESS_TERMINATION_CHECK_COUNTS:
+            # 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')
+
+        self.shared_dict_update_queue.put(dict(self._processes))
+
+    def _update_shared_dict(self, shared_dict_update_queue):
+        while self._running:
+            current_process_pool = shared_dict_update_queue.get()
+            if self.process_pool_shared_dict_proxy is not None:
+                self.process_pool_shared_dict_proxy.clear()
+                self.process_pool_shared_dict_proxy.update(current_process_pool)
+
+
+class PidProcessConverter:
+    """
+    This class is used to get the process details by PID from the process pool shared dict proxy.
+    """
+
+    def __init__(
+            self,
+            process_pool_shared_dict_proxy: multiprocessing.managers.DictProxy
+    ):
+        """
+        :param process_pool_shared_dict_proxy: multiprocessing.managers.DictProxy, multiprocessing shared dict proxy.
+        """
+
+        self.process_pool_shared_dict_proxy: multiprocessing.managers.DictProxy = process_pool_shared_dict_proxy
+
+        self.get_process_with_psutil = get_process_list.GetProcessList(get_method='psutil', connect_on_init=True)
+
+    def get_process_by_pid(self, pid: int):
+        """
+        THIS FUNCTION WILL RUN OUTSIDE THE PROCESS POOL PROCESS. Inside the function that needs
+        the pid to process conversion.
+        Get the process details by PID from the process pool shared dict proxy.
+
+        :param pid: int, the process ID.
+        :return: dict, the process details.
+        """
+
+        counter = 0
+        process_dict: dict = dict()
+        while counter < WAIT_FOR_PROCESS_POLLER_PID_COUNTS:
+            # We need it so that the pool will not change in the middle of the process.
+            current_pid_pool = convert_proxy_dict_to_dict(self.process_pool_shared_dict_proxy)
+            if pid not in current_pid_pool:
+                # print(dict(self.process_pool_shared_dict_proxy))
+                time.sleep(0.1)
+                counter += 1
+            else:
+                process_dict = current_pid_pool[pid]
+                break
+
+        if not process_dict:
+            print_api(f"Error: The PID [{pid}] is not in the pool, trying psutil snapshot.", color='yellow')
+            # Last resort, try to get the process name by current process snapshot.
+            processes = self.get_process_with_psutil.get_processes(as_dict=True)
+            if pid not in processes:
+                print_api(f"Error: Couldn't get the process name for PID: {pid}.", color='red')
+                process_dict = {
+                    'name': pid,
+                    'cmdline': ''
+                }
+            else:
+                process_dict = processes[pid]
+
+        return process_dict
+
+
+def convert_proxy_dict_to_dict(proxy_dict: multiprocessing.managers.DictProxy) -> dict:
+    """
+    Convert the multiprocessing shared dict proxy to a normal dict.
+
+    :param proxy_dict: multiprocessing.managers.DictProxy, the shared dict proxy.
+    :return: dict, the normal dict.
+    """
+
+    # Create a snapshot of the keys
+    keys = list(proxy_dict.keys())
+    current_pid_pool = {}
+
+    for key in keys:
+        try:
+            # Attempt to retrieve the value for each key
+            current_pid_pool[key] = proxy_dict[key]
+        except KeyError:
+            # The key was removed concurrently; skip it
+            continue
+
+    return current_pid_pool

atomicshop/process_poller/tracer_base.py

@@ -0,0 +1,45 @@
+import threading
+
+from .. import get_process_list
+
+
+class Tracer:
+    """
+    Main tracer class for getting the list of processes.
+    """
+
+    def __init__(self):
+        self.tracer_instance = None
+        self.process_queue = None
+
+        self._cycle_callable = None
+        self._processes = {}
+
+    def start(self):
+        """
+        Start the tracer.
+        """
+
+        self._processes = get_process_list.GetProcessList(
+            get_method='pywin32', connect_on_init=True).get_processes(as_dict=True)
+
+        self.process_queue.put(self._processes)
+
+        self.tracer_instance.start()
+
+        thread = threading.Thread(target=self.update_queue)
+        thread.daemon = True
+        thread.start()
+
+    def update_queue(self):
+        """
+        Update the list of processes.
+        """
+
+        while True:
+            # Get subclass specific get process cycle.
+            current_processes = self._cycle_callable()
+
+            self._processes.update(current_processes)
+
+            self.process_queue.put(self._processes)

atomicshop/process_poller/tracers/event_log.py

@@ -0,0 +1,46 @@
+from pathlib import Path
+
+from ...wrappers.pywin32w.win_event_log.subscribes import process_create, process_terminate
+from .. import tracer_base
+
+
+class TracerEventlog(tracer_base.Tracer):
+    """
+    Tracer subclass for getting the list of processes with Windows Event Log Subscription.
+    """
+    def __init__(
+            self,
+            process_queue
+    ):
+        """
+        :param process_queue: Queue. The queue to put the processes in. If None, the processes will not be put in the
+            queue.
+        """
+        super().__init__()
+
+        self.tracer_instance = process_create.ProcessCreateSubscriber()
+
+        self.process_queue = process_queue
+
+        # This function is used in the thread start in the main Tracer class, 'start' function.
+        self._cycle_callable = self.emit_cycle
+
+    def start(self):
+        """
+        Start the tracer.
+        """
+        super().start()
+
+    def emit_cycle(self):
+        """
+        Get the list of processes.
+        """
+
+        # Get the current processes and reinitialize the instance of the dict.
+        current_cycle: dict = self.tracer_instance.emit()
+        current_processes: dict = {int(current_cycle['NewProcessIdInt']): {
+            'name': Path(current_cycle['NewProcessName']).name,
+            'cmdline': current_cycle['CommandLine']}
+        }
+
+        return current_processes

atomicshop/process_poller/tracers/sysmon_etw.py

@@ -0,0 +1,68 @@
+from .. import tracer_base
+from ...etws.traces import trace_sysmon_process_creation
+
+
+class TracerSysmonEtw(tracer_base.Tracer):
+    """
+    Tracer subclass for getting the list of processes with SysMon ETW.
+    """
+    def __init__(
+            self,
+            attrs: list = None,
+            settings: dict = None,
+            process_queue=None
+    ):
+        """
+        :param attrs: list. Default is ['pid', 'original_file_name', 'command_line'].
+            The list of attributes to get from the sysmon trace output. Check SysmonProcessCreationTrace class for
+            available attributes.
+        :settings: dict. The settings dictionary:
+            'sysmon_etw_session_name': str. The Sysmon ETW session name. If None, the default from
+                'trace_sysmon_process_creation' will be used.
+            'sysmon_directory': str. The directory where Sysmon.exe resides. If None, the default from
+                'trace_sysmon_process_creation' will be used.
+        :param process_queue: Queue. The queue to put the processes in. If None, the processes will not be put in the
+            queue.
+        """
+        super().__init__()
+
+        if not attrs:
+            attrs = ['ProcessId', 'OriginalFileName', 'CommandLine']
+
+        if not settings:
+            settings = {
+                'sysmon_etw_session_name': None,
+                'sysmon_directory': None
+            }
+
+        self.tracer_instance = trace_sysmon_process_creation.SysmonProcessCreationTrace(
+            attrs=attrs,
+            session_name=settings.get('sysmon_etw_session_name'),
+            close_existing_session_name=True,
+            sysmon_directory=settings.get('sysmon_directory')
+        )
+
+        self.process_queue = process_queue
+
+        # This function is used in the thread start in the main Tracer class, 'start' function.
+        self._cycle_callable = self.emit_cycle
+
+    def start(self):
+        """
+        Start the tracer.
+        """
+        super().start()
+
+    def emit_cycle(self):
+        """
+        Get the list of processes.
+        """
+
+        # Get the current processes and reinitialize the instance of the dict.
+        current_cycle: dict = self.tracer_instance.emit()
+        current_processes: dict = {int(current_cycle['ProcessId']): {
+            'name': current_cycle['OriginalFileName'],
+            'cmdline': current_cycle['CommandLine']}
+        }
+
+        return current_processes

atomicshop/python_file_patcher.py

@@ -113,4 +113,4 @@ def find_and_replace_in_file(
             for index in found_string_indexes:
                 file_data[index] = file_data[index].replace(single_find.find_what, single_find.replace_to)
 
-    file_io.write_file(content=file_data, file_path=file_path, encoding=encoding, convert_list_to_string=True)
+    file_io.write_file(content=file_data, file_path=file_path, encoding=encoding)

atomicshop/python_functions.py

@@ -1,9 +1,10 @@
 import sys
+from typing import Union
 
 from .print_api import print_api
 
 
-def get_current_python_version_string() -> str:
+def get_python_version_string() -> str:
     """
     Function gets version MAJOR.MINOR.MICRO from 'sys.version_info' object and returns it as a string.
     :return: python MAJOR.MINOR.MICRO version string.
@@ -12,84 +13,35 @@ def get_current_python_version_string() -> str:
     return f'{sys.version_info[0]}.{sys.version_info[1]}.{sys.version_info[2]}'
 
 
-# noinspection PyUnusedLocal
-def check_if_version_object_is_tuple_or_string(version_object: any,
-                                               **kwargs) -> tuple:
-    """
-    Function checks if 'version_object' that was passed is tuple or string.
-    If it's tuple then returns it back. If it's string, converts to tuple then returns it.
-    If the object is none of the above, returns None.
-
-    :param version_object: Can be string ('3.10') or tuple of integers ((3, 10)).
-    :return:
-    """
-    # Check if tuple was passed.
-    if isinstance(version_object, tuple):
-        return version_object
-    else:
-        # Then check if a string was passed.
-        if isinstance(version_object, str):
-            # The check will be against tuple of integers, so we'll convert a string to tuple of integers.
-            return tuple(map(int, version_object.split('.')))
-        else:
-            message = f'[*] Function: [check_if_version_object_is_tuple_or_string]\n' \
-                      f'[*] [version_object] object passed is not tuple or string.\n' \
-                      f'[*] Object type: {type(version_object)}\n' \
-                      f'[*] Object content {version_object}\n' \
-                      f'Exiting...'
-            print_api(message, error_type=True, logger_method='critical', **kwargs)
-
-            return None
-
-
-# noinspection PyUnusedLocal
-def check_python_version_compliance(minimum_version: any,
-                                    maximum_version: any = None,
-                                    **kwargs) -> bool:
+def check_python_version_compliance(
+        min_ver: tuple = None,
+        max_ver: tuple = None,
+) -> str | None:
     """
     Python version check. Should be executed before importing external libraries, since they depend on Python version.
 
-    :param minimum_version: Can be string ('3.10') or tuple of integers ((3, 10)).
-    :param maximum_version: Can be string ('3.10') or tuple of integers ((3, 10)).
+    :param min_ver: tuple of integers (3, 10).
+    :param max_ver: tuple of integers (3, 10).
         If maximum version is not specified, it will be considered as all versions above the minimum are compliant.
-    :return:
+    :return: If version is not compliant, returns string with error message. Otherwise, returns None.
     """
 
-    # Check 'minimum_version' object for string or tuple and get the tuple.
-    minimum_version_scheme: tuple = check_if_version_object_is_tuple_or_string(minimum_version, **kwargs)
-    # If 'maximum_version' object was passed, check it for string or tuple and get the tuple.
-    if maximum_version:
-        maximum_version_scheme: tuple = check_if_version_object_is_tuple_or_string(maximum_version, **kwargs)
-
-    # Get current python version.
-    python_version_full: str = get_current_python_version_string()
-
-    message = f"[*] Current Python Version: {python_version_full}"
-    print_api(message, logger_method='info', **kwargs)
-
-    # if 'maximum_version' passed and current python version is later or equals to the minimum and earlier than
-    # maximum version required.
-    if maximum_version:
-        if not sys.version_info >= minimum_version_scheme or not sys.version_info < maximum_version_scheme:
-            message = f"[!!!] YOU NEED TO INSTALL AT LEAST PYTHON " \
-                      f"{'.'.join(str(i) for i in minimum_version_scheme)}, " \
-                      f"AND EARLIER THAN {'.'.join(str(i) for i in maximum_version_scheme)}, " \
-                      f"to work properly. Unhandled exceptions are inevitable!"
-            print_api(message, error_type=True, logger_method='critical', **kwargs)
-
-            return False
-    # If 'maximum_version' wasn't passed.
+    if not min_ver and not max_ver:
+        raise ValueError("At least one of the version parameters should be passed.")
+
+    current_version_info: tuple = sys.version_info[:3]
+    if min_ver and not max_ver:
+        if current_version_info < min_ver:
+            return f'Python version {".".join(map(str, min_ver))} or higher is required. '\
+                   f'Current version is {".".join(map(str, current_version_info))}.'
+    elif max_ver and not min_ver:
+        if current_version_info > max_ver:
+            return f'Python version up to {".".join(map(str, max_ver))} is required. '\
+                   f'Current version is {".".join(map(str, current_version_info))}.'
+    elif min_ver and max_ver:
+        if not (min_ver <= current_version_info <= max_ver):
+            return f'Python version between {".".join(map(str, min_ver))} and '\
+                   f'{".".join(map(str, max_ver))} is required. '\
+                   f'Current version is {".".join(map(str, current_version_info))}.'
     else:
-        # Check if current python version is later or equals to the minimum required version.
-        if not sys.version_info >= minimum_version_scheme:
-            message = f"[!!!] YOU NEED TO INSTALL AT LEAST PYTHON " \
-                      f"{'.'.join(str(i) for i in minimum_version_scheme)}, " \
-                      f"to work properly. Unhandled exceptions are inevitable!"
-            print_api(message, error_type=True, logger_method='critical', **kwargs)
-
-            return False
-
-    message = "[*] Version Check PASSED."
-    print_api(message, logger_method='info', **kwargs)
-
-    return True
+        return None

atomicshop/question_answer_engine.py

@@ -4,7 +4,7 @@ import sys
 import json
 
 # Custom class imports.
-from atomicshop.filesystem import check_file_existence
+from atomicshop import filesystem
 
 
 class QAEngine:
@@ -20,7 +20,7 @@ class QAEngine:
         # Get 'qa.json' full path.
         qa_fullpath: str = script_directory + os.sep + self.qa_filename
         # Check if it exists.
-        if not check_file_existence(qa_fullpath):
+        if not filesystem.is_file_exists(qa_fullpath):
             print(f'File non-existent: {qa_fullpath}')
             sys.exit()