atomicshop 2.14.2__py3-none-any.whl → 2.14.4__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.2'
+__version__ = '2.14.4'

atomicshop/datetimes.py

@@ -47,7 +47,7 @@ class MonthToNumber:
         'דצמבר': '12'}
 
 
-def get_datetime_from_complex_string_by_pattern(complex_string: str, date_pattern: str):
+def get_datetime_from_complex_string_by_pattern(complex_string: str, date_pattern: str) -> tuple[datetime, str, float]:
     """
     Function will get datetime object from a complex string by pattern.
 
@@ -65,7 +65,8 @@ def get_datetime_from_complex_string_by_pattern(complex_string: str, date_patter
     if date_str:
         # Convert the date string to a datetime object based on the given pattern
         date_obj = datetime.datetime.strptime(date_str.group(), date_pattern)
-        return date_obj
+        date_timestamp = date_obj.timestamp()
+        return date_obj, date_str.group(), date_timestamp
     else:
         raise ValueError("No valid date found in the string")
 

atomicshop/etws/providers.py

@@ -1,5 +1,21 @@
+from typing import Literal
+
 from ..wrappers.ctyping.etw_winapi import etw_functions
 
 
-def get_providers():
-    return etw_functions.get_all_providers()
+def get_providers(key_as: Literal['name', 'guid'] = 'name'):
+    return etw_functions.get_all_providers(key_as=key_as)
+
+
+def get_provider_guid_by_name(provider_name):
+    providers = get_providers(key_as='name')
+
+    try:
+        provider_guid = providers[provider_name]
+    except KeyError:
+        provider_guid = None
+
+    if not provider_guid:
+        raise ValueError(f"Provider '{provider_name}' not found")
+
+    return provider_guid

atomicshop/etws/sessions.py

@@ -1,7 +1,7 @@
 from ..wrappers.ctyping.etw_winapi import etw_functions
 
 
-def stop_and_delete(session_name) -> tuple[bool, int]:
+def stop_and_delete(session_name: str) -> tuple[bool, int]:
     """
     Stop and delete ETW session.
 

atomicshop/etws/trace.py

@@ -1,19 +1,16 @@
 import queue
 import sys
 import time
-from typing import Literal
 
 # Import FireEye Event Tracing library.
 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 +23,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 +37,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 +73,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 +135,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/file_io/csvs.py

@@ -36,7 +36,7 @@ def read_csv_to_list_of_dicts_by_header(
             All the lines of the CSV file will be considered as content.
     :param file_object: file object of the 'open()' function in the decorator. Decorator executes the 'with open()'
         statement and passes to this function. That's why the default is 'None', since we get it from the decorator.
-    :return: list.
+    :return: tuple(list of entries, header(list of cell names)).
     """
 
     # The header fields will be separated to list of "csv_reader.fieldnames".

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}')