atomicshop 2.13.1__py3-none-any.whl → 2.13.2__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.13.1'
+__version__ = '2.13.2'

atomicshop/etws/trace.py

@@ -1,11 +1,20 @@
 import queue
 import sys
+import time
 
 # Import FireEye Event Tracing library.
 import etw
 
 from ..print_api import print_api
 from . import sessions
+from .. import process_poller
+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
 
 
 class EventTrace(etw.ETW):
@@ -15,7 +24,9 @@ class EventTrace(etw.ETW):
             event_callback=None,
             event_id_filters: list = None,
             session_name: str = None,
-            close_existing_session_name: bool = True
+            close_existing_session_name: bool = True,
+            enable_process_poller: bool = False,
+            process_poller_etw_session_name: str = None
     ):
         """
         :param providers: List of tuples with provider name and provider GUID.
@@ -26,7 +37,13 @@ class EventTrace(etw.ETW):
             The default in the 'etw.ETW' method is 'None'.
         :param session_name: The name of the session to create. If not provided, a UUID will be generated.
         :param close_existing_session_name: Boolean to close existing session names.
+        :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_etw_session_name: The name of the ETW session for tracing process creation.
+
         ------------------------------------------
+
         You should stop the ETW tracing when you are done with it.
             'pywintrace' module starts a new session for ETW tracing, and it will not stop the session when the script
             exits or exception is raised.
@@ -46,6 +63,8 @@ class EventTrace(etw.ETW):
         """
         self.event_queue = queue.Queue()
         self.close_existing_session_name: bool = close_existing_session_name
+        self.enable_process_poller: bool = enable_process_poller
+        self.process_poller_etw_session_name: str = process_poller_etw_session_name
 
         # If no callback function is provided, we will use the default one, which will put the event in the queue.
         if not event_callback:
@@ -59,12 +78,23 @@ class EventTrace(etw.ETW):
         for provider in providers:
             etw_format_providers.append(etw.ProviderInfo(provider[0], etw.GUID(provider[1])))
 
+        if not process_poller_etw_session_name:
+            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='sysmon_etw',
+                sysmon_etw_session_name=process_poller_etw_session_name)
+
         super().__init__(
             providers=etw_format_providers, event_callback=function_callable, event_id_filters=event_id_filters,
             session_name=session_name
         )
 
     def start(self):
+        if self.enable_process_poller:
+            self.process_poller.start()
+
         # Check if the session name already exists.
         if sessions.is_session_running(self.session_name):
             print_api(f'ETW Session already running: {self.session_name}', color='yellow')
@@ -87,6 +117,9 @@ class EventTrace(etw.ETW):
     def stop(self):
         super().stop()
 
+        if self.enable_process_poller:
+            self.process_poller.stop()
+
     def emit(self):
         """
         The Function will return the next event from the queue.
@@ -104,34 +137,44 @@ class EventTrace(etw.ETW):
         :return: etw event object.
         """
 
-        return self.event_queue.get()
-
-
-def find_sessions_by_provider(provider_name: str):
-    """
-    Find ETW session by provider name.
-
-    :param provider_name: The name of the provider to search for.
-    """
+        # Get the processes first, since we need the process name and command line.
+        # If they're not ready, we will get just pids from DNS tracing.
+        if self.enable_process_poller:
+            self._get_processes_from_poller()
 
-    return
+        event: tuple = self.event_queue.get()
 
+        event_dict: dict = {
+            'event_id': event[0],
+            'event': event[1],
+            'pid': event[1]['EventHeader']['ProcessId']
+        }
 
-def get_all_providers_from_session(session_name: str):
-    """
-    Get all providers that ETW session uses.
+        if self.enable_process_poller:
+            processes = self.process_poller.get_processes()
+            if event_dict['pid'] not in processes:
+                counter = 0
+                while counter < WAIT_FOR_PROCESS_POLLER_PID_COUNTS:
+                    processes = self.process_poller.get_processes()
+                    if event_dict['pid'] not in processes:
+                        time.sleep(0.1)
+                        counter += 1
+                    else:
+                        break
 
-    :param session_name: The name of the session to get providers from.
-    """
+                if counter == WAIT_FOR_PROCESS_POLLER_PID_COUNTS:
+                    print_api(f"Error: Couldn't get the process name for PID: {event_dict['pid']}.", color='red')
 
-    return
+            event_dict = psutilw.cross_single_connection_with_processes(event_dict, processes)
 
+        return event_dict
 
-def stop_session_by_name(session_name: str):
-    """
-    Stop ETW session by name.
+    def _get_processes_from_poller(self):
+        processes: dict = {}
+        while not processes:
+            processes = self.process_poller.get_processes()
 
-    :param session_name: The name of the session to stop.
-    """
+            if isinstance(processes, BaseException):
+                raise processes
 
-    return
+        return processes

atomicshop/etws/traces/trace_dns.py

@@ -19,20 +19,15 @@ WAIT_FOR_PROCESS_POLLER_PID_COUNTS: int = WAIT_FOR_PROCESS_POLLER_PID_SECONDS *
 
 
 class DnsRequestResponseTrace:
+    """DnsTrace class use to trace DNS events from Windows Event Tracing for EventId 3008."""
     def __init__(
             self,
             attrs: list = None,
             session_name: str = None,
             close_existing_session_name: bool = True,
-            enable_process_poller: bool = False,
             process_poller_etw_session_name: str = None
     ):
         """
-        DnsTrace class use to trace DNS events from Windows Event Tracing for EventId 3008.
-
-        :param enable_process_poller: Boolean to enable process poller. Gets the process PID, Name and CommandLine,
-            every 100 ms. 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 attrs: List of attributes to return. If None, all attributes will be returned.
         :param session_name: The name of the session to create. If not provided, a UUID will be generated.
         :param close_existing_session_name: Boolean to close existing session names.
@@ -63,7 +58,6 @@ class DnsRequestResponseTrace:
             dns_trace_w.stop()
         """
 
-        self.enable_process_poller = enable_process_poller
         self.attrs = attrs
 
         if not session_name:
@@ -74,26 +68,17 @@ class DnsRequestResponseTrace:
             # lambda x: self.event_queue.put(x),
             event_id_filters=[REQUEST_RESP_EVENT_ID],
             session_name=session_name,
-            close_existing_session_name=close_existing_session_name
+            close_existing_session_name=close_existing_session_name,
+            enable_process_poller=True,
+            process_poller_etw_session_name=process_poller_etw_session_name
         )
 
-        if self.enable_process_poller:
-            self.process_poller = ProcessPollerPool(
-                operation='process', poller_method='sysmon_etw',
-                sysmon_etw_session_name=process_poller_etw_session_name)
-
     def start(self):
-        if self.enable_process_poller:
-            self.process_poller.start()
-
         self.event_trace.start()
 
     def stop(self):
         self.event_trace.stop()
 
-        if self.enable_process_poller:
-            self.process_poller.stop()
-
     def emit(self):
         """
         Function that will return the next event from the queue.
@@ -107,29 +92,26 @@ class DnsRequestResponseTrace:
         :return: Dictionary with the event data.
         """
 
-        # Get the processes first, since we need the process name and command line.
-        # If they're not ready, we will get just pids from DNS tracing.
-        if self.enable_process_poller:
-            self._get_processes_from_poller()
-
         event = self.event_trace.emit()
 
         event_dict: dict = {
-            'pid': event[1]['EventHeader']['ProcessId'],
-            'etw_id': event[0],
-            'domain': event[1]['QueryName'],
-            'query_type_id': str(event[1]['QueryType']),
-            'query_type': dns.TYPES_DICT[str(event[1]['QueryType'])]
+            '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'])],
+            'pid': event['pid'],
+            'name': event['name'],
+            'cmdline': event['cmdline']
         }
 
         # 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[1]['QueryResults']:
+        if event['event']['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[1]['QueryResults'].split(';')
+            list_of_parameters = event['event']['QueryResults'].split(';')
 
             # Iterating through all the parameters that we got from 'QueryResults' key.
             for parameter in list_of_parameters:
@@ -149,47 +131,17 @@ class DnsRequestResponseTrace:
         event_dict['other_domains'] = list_of_other_domains
 
         # Getting the 'QueryStatus' key.
-        event_dict['status_id'] = event[1]['QueryStatus']
+        event_dict['status_id'] = event['event']['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[1]['QueryStatus'] == '0':
+        if event['event']['QueryStatus'] == '0':
             event_dict['status'] = 'Success'
         else:
             event_dict['status'] = 'Error'
 
-        if self.enable_process_poller:
-            processes = self.process_poller.get_processes()
-            if event_dict['pid'] not in processes:
-                counter = 0
-                while counter < WAIT_FOR_PROCESS_POLLER_PID_COUNTS:
-                    processes = self.process_poller.get_processes()
-                    if event_dict['pid'] not in processes:
-                        time.sleep(0.1)
-                        counter += 1
-                    else:
-                        break
-
-                if counter == WAIT_FOR_PROCESS_POLLER_PID_COUNTS:
-                    print_api(f"Error: Couldn't get the process name for PID: {event_dict['pid']}.", color='red')
-
-            event_dict = psutilw.cross_single_connection_with_processes(event_dict, processes)
-            # If it was impossible to get the process name from the process poller, get it from psutil.
-            # if event_dict['name'].isnumeric():
-            #     event_dict['name'] = process_name
-
         if self.attrs:
             event_dict = dicts.reorder_keys(
                 event_dict, self.attrs, skip_keys_not_in_list=True)
 
         return event_dict
-
-    def _get_processes_from_poller(self):
-        processes: dict = {}
-        while not processes:
-            processes = self.process_poller.get_processes()
-
-            if isinstance(processes, BaseException):
-                raise processes
-
-        return processes

atomicshop/etws/traces/trace_sysmon_process_creation.py

@@ -1,8 +1,6 @@
 from .. import trace, const
-from ...wrappers.psutilw import psutilw
 from ...wrappers import sysmonw
 from ...basics import dicts
-from ...print_api import print_api
 
 
 PROVIDER_NAME: str = const.ETW_SYSMON['provider_name']
@@ -87,26 +85,26 @@ class SysmonProcessCreationTrace:
         event = self.event_trace.emit()
 
         event_dict: dict = {
-            'etw_id': event[0],
-            'pid': event[1]['ProcessId'],
-            'process_guid': event[1]['ProcessGuid'],
-            'image': event[1]['Image'],
-            'file_version': event[1]['FileVersion'],
-            'product': event[1]['Product'],
-            'company': event[1]['Company'],
-            'original_file_name': event[1]['OriginalFileName'],
-            'command_line': event[1]['CommandLine'],
-            'current_directory': event[1]['CurrentDirectory'],
-            'user': event[1]['User'],
-            'logon_id': event[1]['LogonId'],
-            'logon_guid': event[1]['LogonGuid'],
-            'terminal_session_id': event[1]['TerminalSessionId'],
-            'integrity_level': event[1]['IntegrityLevel'],
-            'hashes': event[1]['Hashes'],
-            'parent_process_guid': event[1]['ParentProcessGuid'],
-            'parent_process_id': event[1]['ParentProcessId'],
-            'parent_image': event[1]['ParentImage'],
-            'parent_command_line': event[1]['ParentCommandLine']
+            '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']
         }
 
         if self.attrs:

atomicshop/monitor/checks/dns.py

@@ -28,7 +28,6 @@ class DnsCheck:
                 attrs=['name', 'cmdline', 'domain', 'query_type'],
                 session_name=self.etw_session_name,
                 close_existing_session_name=True,
-                enable_process_poller=True,
                 process_poller_etw_session_name=change_monitor_instance.etw_process_session_name
             )
         )

atomicshop/process_poller.py

@@ -6,7 +6,7 @@ from typing import Literal, Union
 from .wrappers.pywin32w import wmi_win32process
 from .wrappers.psutilw import psutilw
 from .etws.traces import trace_sysmon_process_creation
-from .basics import list_of_dicts, dicts
+from .basics import dicts
 from .process_name_cmd import ProcessNameCmdline
 from .print_api import print_api
 
@@ -322,6 +322,7 @@ def _worker(
             running_state = False
             exception = e
             print_api(f'Exception in ProcessPollerPool: {e}', color='red')
+            raise
 
     if not running_state:
         process_queue.put(exception)

{atomicshop-2.13.1.dist-info → atomicshop-2.13.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: atomicshop
-Version: 2.13.1
+Version: 2.13.2
 Summary: Atomic functions and classes to make developer life easier
 Author: Denis Kras
 License: MIT License

{atomicshop-2.13.1.dist-info → atomicshop-2.13.2.dist-info}/RECORD

@@ -1,4 +1,4 @@
-atomicshop/__init__.py,sha256=aqwPLStJUJ8YnGq3I420Yccvy8q9JIXVsivcavrcOtM,123
+atomicshop/__init__.py,sha256=e_cepUB_kyEwwtls5dudclY5653g8Bn7epJj7a_khw8,123
 atomicshop/_basics_temp.py,sha256=6cu2dd6r2dLrd1BRNcVDKTHlsHs_26Gpw8QS6v32lQ0,3699
 atomicshop/_create_pdf_demo.py,sha256=Yi-PGZuMg0RKvQmLqVeLIZYadqEZwUm-4A9JxBl_vYA,3713
 atomicshop/_patch_import.py,sha256=ENp55sKVJ0e6-4lBvZnpz9PQCt3Otbur7F6aXDlyje4,6334
@@ -27,7 +27,7 @@ atomicshop/permissions.py,sha256=P6tiUKV-Gw-c3ePEVsst9bqWaHJbB4ZlJB4xbDYVpEs,443
 atomicshop/print_api.py,sha256=DhbCQd0MWZZ5GYEk4oTu1opRFC-b31g1VWZgTGewG2Y,11568
 atomicshop/process.py,sha256=R1BtXWjG2g2Q3WlsyhbIlXZz0UkQeagY7fQyBOIX_DM,15951
 atomicshop/process_name_cmd.py,sha256=CtaSp3mgxxJKCCVW8BLx6BJNx4giCklU_T7USiCEwfc,5162
-atomicshop/process_poller.py,sha256=sEuuFu3gmUrp4tY1LVxE27L0eK5LIwNKjTbV7KtzeRM,15029
+atomicshop/process_poller.py,sha256=sGIiLNdYdyxdER7hKP8QcbYg1esKe5ymCfh5qUATpJM,15033
 atomicshop/python_file_patcher.py,sha256=kd3rBWvTcosLEk-7TycNdfKW9fZbe161iVwmH4niUo0,5515
 atomicshop/python_functions.py,sha256=zJg4ogUwECxrDD7xdDN5JikIUctITM5lsyabr_ZNsRw,4435
 atomicshop/question_answer_engine.py,sha256=DuOn7QEgKKfqZu2cR8mVeFIfFgayfBHiW-jY2VPq_Fo,841
@@ -106,10 +106,10 @@ atomicshop/etws/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 atomicshop/etws/const.py,sha256=v3x_IdCYeSKbCGywiZFOZln80ldpwKW5nuMDuUe51Jg,1257
 atomicshop/etws/providers.py,sha256=fVmWi-uGdtnsQTDpu_ty6dzx0GMhGokiST73LNBEJ38,129
 atomicshop/etws/sessions.py,sha256=k3miewU278xn829cqDbsuH_bmZHPQE9-Zn-hINbxUSE,1330
-atomicshop/etws/trace.py,sha256=tjBvV4RxCSn8Aoxj8NVxmqHekdECne_y4Zv2lbh11ak,5341
+atomicshop/etws/trace.py,sha256=ptdDyLuk8frKgmgOEug8KrH5niC6SfgWnOqw06tGPAw,7831
 atomicshop/etws/traces/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-atomicshop/etws/traces/trace_dns.py,sha256=wRFwfFgVgMbYJu1Sf0Yy9LPoMrf9kXejl2Xjk3PNELQ,8430
-atomicshop/etws/traces/trace_sysmon_process_creation.py,sha256=pjb0qzNZttucEWOKv2j_BGVqaSeKjBL4O8oCsAteiGU,4785
+atomicshop/etws/traces/trace_dns.py,sha256=8Fl5UW4An3SE8tGXGTufTgOcWa574YZ6zQbU7x5DYrs,6270
+atomicshop/etws/traces/trace_sysmon_process_creation.py,sha256=WdlQiOfRZC-_PpuE6BkOw5zB0DLc3fhVrANyLrgfCac,4833
 atomicshop/file_io/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 atomicshop/file_io/csvs.py,sha256=y8cJtnlN-NNxNupzJgSeGq9aQ4wNxYLFPX9vNNlUiIc,5830
 atomicshop/file_io/docxs.py,sha256=6tcYFGp0vRsHR47VwcRqwhdt2DQOwrAUYhrwN996n9U,5117
@@ -141,7 +141,7 @@ atomicshop/mitm/engines/__reference_general/responder___reference_general.py,sha
 atomicshop/monitor/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 atomicshop/monitor/change_monitor.py,sha256=eJRP7NBnA-Y_n6Ofm_jKrR5XguJV6gEmxFdtdGMBF3k,7866
 atomicshop/monitor/checks/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-atomicshop/monitor/checks/dns.py,sha256=_DxInUSiiMa8Bx1sYJ6fRzRl7n_eAsoaGHhUk9XgF7w,7182
+atomicshop/monitor/checks/dns.py,sha256=OJWvMKl0zNtvBTYVYHSK_tfH0cNxCJ1kc_O-2Ob_5fc,7137
 atomicshop/monitor/checks/file.py,sha256=2tIDSlX2KZNc_9i9ji1tcOqupbFTIOj7cKXLyBEDWMk,3263
 atomicshop/monitor/checks/network.py,sha256=CGZWl4WlQrxayZeVF9JspJXwYA-zWx8ECWTVGSlXc98,3825
 atomicshop/monitor/checks/process_running.py,sha256=x66wd6-l466r8sbRQaIli0yswyGt1dH2DVXkGDL6O0Q,1891
@@ -265,8 +265,8 @@ atomicshop/wrappers/socketw/socket_server_tester.py,sha256=AhpurHJmP2kgzHaUbq5ey
 atomicshop/wrappers/socketw/socket_wrapper.py,sha256=aXBwlEIJhFT0-c4i8iNlFx2It9VpCEpsv--5Oqcpxao,11624
 atomicshop/wrappers/socketw/ssl_base.py,sha256=k4V3gwkbq10MvOH4btU4onLX2GNOsSfUAdcHmL1rpVE,2274
 atomicshop/wrappers/socketw/statistics_csv.py,sha256=t3dtDEfN47CfYVi0CW6Kc2QHTEeZVyYhc57IYYh5nmA,826
-atomicshop-2.13.1.dist-info/LICENSE.txt,sha256=lLU7EYycfYcK2NR_1gfnhnRC8b8ccOTElACYplgZN88,1094
-atomicshop-2.13.1.dist-info/METADATA,sha256=yOVvRROyXl4BM9Qr9EzI5BTZgvEo0OG3-uCo9Ryni2Y,10478
-atomicshop-2.13.1.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-atomicshop-2.13.1.dist-info/top_level.txt,sha256=EgKJB-7xcrAPeqTRF2laD_Np2gNGYkJkd4OyXqpJphA,11
-atomicshop-2.13.1.dist-info/RECORD,,
+atomicshop-2.13.2.dist-info/LICENSE.txt,sha256=lLU7EYycfYcK2NR_1gfnhnRC8b8ccOTElACYplgZN88,1094
+atomicshop-2.13.2.dist-info/METADATA,sha256=FwrMOVELwpnXQKf-LbaJhMr9B2mbEos5YA94BQbSuAk,10478
+atomicshop-2.13.2.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+atomicshop-2.13.2.dist-info/top_level.txt,sha256=EgKJB-7xcrAPeqTRF2laD_Np2gNGYkJkd4OyXqpJphA,11
+atomicshop-2.13.2.dist-info/RECORD,,