atomicshop 2.14.1__py3-none-any.whl → 2.14.3__py3-none-any.whl

atomicshop/__init__.py

@@ -1,4 +1,4 @@
 """Atomic Basic functions and classes to make developer life easier"""
 
 __author__ = "Den Kras"
-__version__ = '2.14.1'
+__version__ = '2.14.3'

atomicshop/etws/trace.py

@@ -8,12 +8,10 @@ import etw
 
 from ..print_api import print_api
 from . import sessions
-from .. import process_poller
+from ..process_poller import simple_process_pool
 from ..wrappers.psutilw import psutilw
 
 
-PROCESS_POLLER_ETW_DEFAULT_SESSION_NAME: str = 'AtomicShopProcessTrace'
-
 WAIT_FOR_PROCESS_POLLER_PID_SECONDS: int = 3
 WAIT_FOR_PROCESS_POLLER_PID_COUNTS: int = WAIT_FOR_PROCESS_POLLER_PID_SECONDS * 10
 
@@ -26,8 +24,7 @@ class EventTrace(etw.ETW):
             event_id_filters: list = None,
             session_name: str = None,
             close_existing_session_name: bool = True,
-            enable_process_poller: bool = False,
-            process_poller_method: Literal['psutil', 'pywin32', 'process_dll', 'sysmon_etw', 'event_log'] = 'event_log'
+            enable_process_poller: bool = False
     ):
         """
         :param providers: List of tuples with provider name and provider GUID.
@@ -41,13 +38,6 @@ class EventTrace(etw.ETW):
         :param enable_process_poller: Boolean to enable process poller. Gets the process PID, Name and CommandLine.
             Since the DNS events doesn't contain the process name and command line, only PID.
             Then DNS events will be enriched with the process name and command line from the process poller.
-        :param process_poller_method: The method to get the process information. For more information, see the
-            'process_poller.ProcessPollerPool' class. Summary:
-                'psutil': Uses 'psutil' library to get the process information.
-                'pywin32': Uses 'pywin32' library to get the process information.
-                'process_dll': Uses 'process' custom DLL to get the process information.
-                'sysmon_etw': Uses 'sysmon_etw' uses sysmon and ETW to get the process information.
-                'event_log': Uses Security Windows EVent Log channel (event id 4688) to get the process information.
 
         ------------------------------------------
 
@@ -84,14 +74,8 @@ class EventTrace(etw.ETW):
         for provider in providers:
             etw_format_providers.append(etw.ProviderInfo(provider[0], etw.GUID(provider[1])))
 
-        process_poller_etw_session_name = None
-        if process_poller_method == 'sysmon_etw':
-            process_poller_etw_session_name = PROCESS_POLLER_ETW_DEFAULT_SESSION_NAME
-
         if self.enable_process_poller:
-            self.process_poller = process_poller.ProcessPollerPool(
-                operation='process', poller_method=process_poller_method,
-                sysmon_etw_session_name=process_poller_etw_session_name)
+            self.process_poller = simple_process_pool.SimpleProcessPool()
 
         super().__init__(
             providers=etw_format_providers, event_callback=function_callable, event_id_filters=event_id_filters,
@@ -152,8 +136,8 @@ class EventTrace(etw.ETW):
         event: tuple = self.event_queue.get()
 
         event_dict: dict = {
-            'event_id': event[0],
-            'event': event[1],
+            'EventId': event[0],
+            'EventHeader': event[1],
             'pid': event[1]['EventHeader']['ProcessId']
         }
 

atomicshop/etws/traces/trace_dns.py

@@ -1,10 +1,5 @@
-import time
-
 from .. import trace, const
-from ...wrappers.psutilw import psutilw
 from ...basics import dicts
-from ...process_poller import ProcessPollerPool
-from ...print_api import print_api
 from ... import dns
 
 
@@ -24,7 +19,8 @@ class DnsRequestResponseTrace:
             self,
             attrs: list = None,
             session_name: str = None,
-            close_existing_session_name: bool = True
+            close_existing_session_name: bool = True,
+            skip_record_list: list = None
     ):
         """
         :param attrs: List of attributes to return. If None, all attributes will be returned.
@@ -35,6 +31,7 @@ class DnsRequestResponseTrace:
             False: if ETW session with 'session_name' exists, you will be notified and the new session will not be
                 created. Instead, the existing session will be used. If there is a buffer from the previous session,
                 you will get the events from the buffer.
+        :param skip_record_list: List of DNS Records to skip emitting. Example: ['PTR', 'SRV']
 
         -------------------------------------------------
 
@@ -56,6 +53,7 @@ class DnsRequestResponseTrace:
         """
 
         self.attrs = attrs
+        self.skip_record_list = skip_record_list
 
         if not session_name:
             session_name = ETW_DEFAULT_SESSION_NAME
@@ -91,23 +89,28 @@ class DnsRequestResponseTrace:
         event = self.event_trace.emit()
 
         event_dict: dict = {
-            'event_id': event['event_id'],
-            'domain': event['event']['QueryName'],
-            'query_type_id': str(event['event']['QueryType']),
-            'query_type': dns.TYPES_DICT[str(event['event']['QueryType'])],
+            'event_id': event['EventId'],
+            'domain': event['EventHeader']['QueryName'],
+            'query_type_id': str(event['EventHeader']['QueryType']),
+            'query_type': dns.TYPES_DICT[str(event['EventHeader']['QueryType'])],
             'pid': event['pid'],
             'name': event['name'],
             'cmdline': event['cmdline']
         }
 
+        # Skip emitting the record if it is in the 'skip_record_list'.
+        # Just recall the function to get the next event and return it.
+        if event_dict['query_type'] in self.skip_record_list:
+            return self.emit()
+
         # Defining list if ips and other answers, which aren't IPs.
         list_of_ips = list()
         list_of_other_domains = list()
         # Parse DNS results, only if 'QueryResults' key isn't empty, since many of the events are, mostly due errors.
-        if event['event']['QueryResults']:
+        if event['EventHeader']['QueryResults']:
             # 'QueryResults' key contains a string with all the 'Answers' divided by type and ';' character.
             # Basically, we can parse each type out of string, but we need only IPs and other answers.
-            list_of_parameters = event['event']['QueryResults'].split(';')
+            list_of_parameters = event['EventHeader']['QueryResults'].split(';')
 
             # Iterating through all the parameters that we got from 'QueryResults' key.
             for parameter in list_of_parameters:
@@ -127,11 +130,11 @@ class DnsRequestResponseTrace:
         event_dict['other_domains'] = list_of_other_domains
 
         # Getting the 'QueryStatus' key.
-        event_dict['status_id'] = event['event']['QueryStatus']
+        event_dict['status_id'] = event['EventHeader']['QueryStatus']
 
         # Getting the 'QueryStatus' key. If DNS Query Status is '0' then it was executed successfully.
         # And if not, it means there was an error. The 'QueryStatus' indicate what number of an error it is.
-        if event['event']['QueryStatus'] == '0':
+        if event['EventHeader']['QueryStatus'] == '0':
             event_dict['status'] = 'Success'
         else:
             event_dict['status'] = 'Error'

atomicshop/etws/traces/trace_sysmon_process_creation.py

@@ -3,6 +3,8 @@ from ...wrappers import sysmonw
 from ...basics import dicts
 
 
+DEFAULT_SESSION_NAME: str = 'AtomicShopSysmonProcessCreationTrace'
+
 PROVIDER_NAME: str = const.ETW_SYSMON['provider_name']
 PROVIDER_GUID: str = const.ETW_SYSMON['provider_guid']
 PROCESS_CREATION_EVENT_ID: int = const.ETW_SYSMON['event_ids']['process_create']
@@ -53,6 +55,9 @@ class SysmonProcessCreationTrace:
         self.attrs = attrs
         self.sysmon_directory: str = sysmon_directory
 
+        if not session_name:
+            session_name = DEFAULT_SESSION_NAME
+
         self.event_trace = trace.EventTrace(
             providers=[(PROVIDER_NAME, PROVIDER_GUID)],
             # lambda x: self.event_queue.put(x),
@@ -80,32 +85,39 @@ class SysmonProcessCreationTrace:
                 print(dns_dict)
 
         :return: Dictionary with the event data.
+
+        -----------------------------------------------
+
+        Structure of the returned dictionary:
+        {
+            'event_id': int,
+            'ProcessId': int,
+            'ProcessGuid': str,
+            'Image': str,
+            'FileVersion': str,
+            'Product': str,
+            'Company': str,
+            'OriginalFileName': str,
+            'CommandLine': str,
+            'CurrentDirectory': str,
+            'User': str,
+            'LogonId': str,
+            'LogonGuid': str,
+            'TerminalSessionId': int,
+            'IntegrityLevel': str,
+            'Hashes': dict,
+            'ParentProcessGuid': str,
+            'ParentProcessId': int,
+            'ParentImage': str,
+            'ParentCommandLine': str
+        }
+
         """
 
         event = self.event_trace.emit()
 
-        event_dict: dict = {
-            'event_id': event['event_id'],
-            'pid': event['event']['ProcessId'],
-            'process_guid': event['event']['ProcessGuid'],
-            'image': event['event']['Image'],
-            'file_version': event['event']['FileVersion'],
-            'product': event['event']['Product'],
-            'company': event['event']['Company'],
-            'original_file_name': event['event']['OriginalFileName'],
-            'command_line': event['event']['CommandLine'],
-            'current_directory': event['event']['CurrentDirectory'],
-            'user': event['event']['User'],
-            'logon_id': event['event']['LogonId'],
-            'logon_guid': event['event']['LogonGuid'],
-            'terminal_session_id': event['event']['TerminalSessionId'],
-            'integrity_level': event['event']['IntegrityLevel'],
-            'hashes': event['event']['Hashes'],
-            'parent_process_guid': event['event']['ParentProcessGuid'],
-            'parent_process_id': event['event']['ParentProcessId'],
-            'parent_image': event['event']['ParentImage'],
-            'parent_command_line': event['event']['ParentCommandLine']
-        }
+        event_dict = {'EventId': event['EventId']}
+        event_dict.update(event['EventHeader'])
 
         if self.attrs:
             event_dict = dicts.reorder_keys(

atomicshop/get_process_list.py

@@ -0,0 +1,133 @@
+from typing import Union, Literal
+
+from .wrappers.pywin32w import wmi_win32process
+from .wrappers.psutilw import psutilw
+from .basics import dicts
+from . import get_process_name_cmd_dll
+
+
+class GetProcessList:
+    """
+    The class is responsible for getting the list of running processes.
+
+    Example of one time polling with 'pywin32' method:
+    from atomicshop import process_poller
+    process_list: dict = \
+        process_poller.GetProcessList(get_method='pywin32', connect_on_init=True).get_processes(as_dict=True)
+    """
+    def __init__(
+            self,
+            get_method: Literal['psutil', 'pywin32', 'process_dll'] = 'process_dll',
+            connect_on_init: bool = False
+    ):
+        """
+        :param get_method: str, The method to get the list of processes. Default is 'process_list_dll'.
+            'psutil': Get the list of processes by 'psutil' library. Resource intensive and slow.
+            'pywin32': Get the list of processes by 'pywin32' library, using WMI. Not resource intensive, but slow.
+            'process_dll'. Not resource intensive and fast. Probably works only in Windows 10 x64.
+        :param connect_on_init: bool, if True, will connect to the service on init. 'psutil' don't need to connect.
+        """
+        self.get_method = get_method
+        self.process_polling_instance = None
+
+        self.connected = False
+
+        if self.get_method == 'psutil':
+            self.process_polling_instance = psutilw.PsutilProcesses()
+            self.connected = True
+        elif self.get_method == 'pywin32':
+            self.process_polling_instance = wmi_win32process.Pywin32Processes()
+        elif self.get_method == 'process_dll':
+            self.process_polling_instance = get_process_name_cmd_dll.ProcessNameCmdline()
+
+        if connect_on_init:
+            self.connect()
+
+    def connect(self):
+        """
+        Connect to the service. 'psutil' don't need to connect.
+        """
+
+        # If poller method is none of the allowed methods.
+        if self.get_method not in ['psutil', 'pywin32', 'process_dll']:
+            raise ValueError(f"Method '{self.get_method}' is not allowed.")
+
+        # If the service is not connected yet. Since 'psutil' don't need to connect.
+        if not self.connected:
+            if self.get_method == 'pywin32':
+                self.process_polling_instance.connect()
+                self.connected = True
+            elif self.get_method == 'process_dll':
+                self.process_polling_instance.load()
+                self.connected = True
+
+    def get_processes(self, as_dict: bool = True) -> Union[list, dict]:
+        """
+        The function will get the list of opened processes and return it as a list of dicts.
+
+        :return: dict while key is pid or list of dicts, of opened processes (depending on 'as_dict' setting).
+        """
+
+        if as_dict:
+            if self.get_method == 'psutil':
+                return self.process_polling_instance.get_processes_as_dict(
+                    attrs=['pid', 'name', 'cmdline'], cmdline_to_string=True)
+            elif self.get_method == 'pywin32':
+                processes = self.process_polling_instance.get_processes_as_dict(
+                    attrs=['ProcessId', 'Name', 'CommandLine'])
+
+                # Convert the keys from WMI to the keys that are used in 'psutil'.
+                converted_process_dict = dict()
+                for pid, process_info in processes.items():
+                    converted_process_dict[pid] = dicts.convert_key_names(
+                        process_info, {'Name': 'name', 'CommandLine': 'cmdline'})
+
+                return converted_process_dict
+            elif self.get_method == 'process_dll':
+                return self.process_polling_instance.get_process_details(as_dict=True)
+        else:
+            if self.get_method == 'psutil':
+                return self.process_polling_instance.get_processes_as_list_of_dicts(
+                    attrs=['pid', 'name', 'cmdline'], cmdline_to_string=True)
+            elif self.get_method == 'pywin32':
+                processes = self.process_polling_instance.get_processes_as_list_of_dicts(
+                    attrs=['ProcessId', 'Name', 'CommandLine'])
+
+                # Convert the keys from WMI to the keys that are used in 'psutil'.
+                for process_index, process_info in enumerate(processes):
+                    processes[process_index] = dicts.convert_key_names(
+                        process_info, {'ProcessId': 'pid', 'Name': 'name', 'CommandLine': 'cmdline'})
+
+                return processes
+            elif self.get_method == 'process_dll':
+                return self.process_polling_instance.get_process_details(as_dict=as_dict)
+
+
+def get_process_time_tester(
+        get_method: Literal['psutil', 'pywin32', 'process_dll'] = 'process_dll',
+        times_to_test: int = 50
+):
+    """
+    The function will test the time it takes to get the list of processes with different methods and cycles.
+
+    :param get_method: str, The method to get the list of processes. Default is 'process_list_dll'.
+        'psutil': Get the list of processes by 'psutil' library. Resource intensive and slow.
+        'pywin32': Get the list of processes by 'pywin32' library, using WMI. Not resource intensive, but slow.
+        'process_dll'. Not resource intensive and fast. Probably works only in Windows 10 x64
+    :param times_to_test: int, how many times to test the function.
+    """
+
+    import timeit
+
+    setup_code = '''
+from atomicshop import process_poller
+get_process_list = process_poller.GetProcessList(get_method=get_method, connect_on_init=True)
+'''
+
+    test_code = '''
+test = get_process_list.get_processes()
+'''
+
+    # globals need to be specified, otherwise the setup_code won't work with passed variables.
+    times = timeit.timeit(setup=setup_code, stmt=test_code, number=times_to_test, globals=locals())
+    print(f'Execution time: {times}')

atomicshop/monitor/change_monitor.py

@@ -4,6 +4,7 @@ from .checks import dns, network, file, url, process_running
 
 
 DNS__DEFAULT_SETTINGS = {
+    'skip_record_list': [],                              # List of DNS Records to skip emitting. Example: ['PTR', 'SRV']
     'learning_mode_create_unique_entries_list': True,
     'learning_hours': 24,                                   # 0 - the learning will never stop.
     'alert_about_missing_entries_after_learning': False,

atomicshop/monitor/checks/dns.py

@@ -27,7 +27,8 @@ class DnsCheck:
             trace_dns.DnsRequestResponseTrace(
                 attrs=['name', 'cmdline', 'domain', 'query_type'],
                 session_name=self.etw_session_name,
-                close_existing_session_name=True
+                close_existing_session_name=True,
+                skip_record_list=self.settings['skip_record_list'],
             )
         )
 

atomicshop/process.py

@@ -11,7 +11,7 @@ from .basics import strings
 from .wrappers import ubuntu_terminal
 
 if os.name == 'nt':
-    from . import process_poller
+    from . import get_process_list
 
 
 def is_command_exists(cmd: str) -> bool:
@@ -265,8 +265,8 @@ def match_pattern_against_running_processes_cmdlines(
     """
 
     # Get the list of all the currently running processes.
-    get_process_list = process_poller.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='pywin32', 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.

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)