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

atomicshop/etws/trace.py

@@ -1,6 +1,7 @@
 import queue
 import sys
-import time
+import multiprocessing.managers
+from datetime import datetime
 
 # Import FireEye Event Tracing library.
 import etw
@@ -8,27 +9,25 @@ import etw
 from ..print_api import print_api
 from . import sessions
 from ..process_poller import simple_process_pool
-from ..wrappers.psutilw import psutilw
-
-
-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):
     def __init__(
             self,
             providers: list,
-            event_callback=None,
+            event_callback: callable = None,
             event_id_filters: list = None,
             session_name: str = None,
             close_existing_session_name: bool = True,
-            enable_process_poller: bool = False
+            enable_process_poller: bool = False,
+            process_pool_shared_dict_proxy: multiprocessing.managers.DictProxy = None
     ):
         """
         :param providers: List of tuples with provider name and provider GUID.
             tuple[0] = provider name
             tuple[1] = provider GUID
+
+            Example: [('Microsoft-Windows-DNS-Client', '{1c95126e-7ee8-4e23-86b2-6e7e4a5a8e9b}')]
         :param event_callback: Reference to the callable callback function that will be called for each occurring event.
         :param event_id_filters: List of event IDs that we want to filter. If not provided, all events will be returned.
             The default in the 'etw.ETW' method is 'None'.
@@ -37,6 +36,13 @@ 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_pool_shared_dict_proxy: multiprocessing.managers.DictProxy,
+            multiprocessing shared dict proxy that contains current processes.
+            Check the 'atomicshop\process_poller\simple_process_pool.py' SimpleProcessPool class for more information.
+
+            If None, the process poller will create a new shared dict proxy.
+            If provided, then the provided shared dict proxy will be used.
+            Off course valid only if 'enable_process_poller' is True.
 
         ------------------------------------------
 
@@ -60,6 +66,7 @@ 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_pool_shared_dict_proxy: multiprocessing.managers.DictProxy = process_pool_shared_dict_proxy
 
         # If no callback function is provided, we will use the default one, which will put the event in the queue.
         if not event_callback:
@@ -73,8 +80,16 @@ class EventTrace(etw.ETW):
         for provider in providers:
             etw_format_providers.append(etw.ProviderInfo(provider[0], etw.GUID(provider[1])))
 
+        self.self_hosted_poller: bool = False
         if self.enable_process_poller:
-            self.process_poller = simple_process_pool.SimpleProcessPool()
+            if self.process_pool_shared_dict_proxy is None:
+                self.self_hosted_poller = True
+                self.process_poller = simple_process_pool.SimpleProcessPool()
+                self.multiprocessing_manager: multiprocessing.managers.SyncManager = multiprocessing.Manager()
+                self.process_pool_shared_dict_proxy = self.multiprocessing_manager.dict()
+
+            self.pid_process_converter = simple_process_pool.PidProcessConverter(
+                process_pool_shared_dict_proxy=self.process_pool_shared_dict_proxy)
 
         super().__init__(
             providers=etw_format_providers, event_callback=function_callable, event_id_filters=event_id_filters,
@@ -82,7 +97,7 @@ class EventTrace(etw.ETW):
         )
 
     def start(self):
-        if self.enable_process_poller:
+        if self.enable_process_poller and self.self_hosted_poller:
             self.process_poller.start()
 
         # Check if the session name already exists.
@@ -107,9 +122,11 @@ class EventTrace(etw.ETW):
     def stop(self):
         super().stop()
 
-        if self.enable_process_poller:
+        if self.self_hosted_poller:
             self.process_poller.stop()
 
+            self.multiprocessing_manager.shutdown()
+
     def emit(self):
         """
         The Function will return the next event from the queue.
@@ -127,44 +144,25 @@ class EventTrace(etw.ETW):
         :return: etw event object.
         """
 
-        # 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: tuple = self.event_queue.get()
 
+        current_datetime = datetime.now()
+        readable_time = current_datetime.strftime('%Y-%m-%d %H:%M:%S.%f')
+
         event_dict: dict = {
             'EventId': event[0],
             'EventHeader': event[1],
-            'pid': event[1]['EventHeader']['ProcessId']
+            'timestamp': readable_time
         }
 
+        if 'ProcessId' not in event[1]:
+            event_dict['pid'] = event[1]['EventHeader']['ProcessId']
+        else:
+            event_dict['pid'] = event[1]['ProcessId']
+
         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)
+            process_info: dict = self.pid_process_converter.get_process_by_pid(event_dict['pid'])
+            event_dict['name'] = process_info['name']
+            event_dict['cmdline'] = process_info['cmdline']
 
         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_dns.py

@@ -1,6 +1,8 @@
+import multiprocessing.managers
+
 from .. import trace, const
 from ...basics import dicts
-from ... import dns
+from ... import dns, ip_addresses
 
 
 ETW_DEFAULT_SESSION_NAME: str = 'AtomicShopDnsTrace'
@@ -9,9 +11,6 @@ PROVIDER_NAME: str = const.ETW_DNS['provider_name']
 PROVIDER_GUID: str = const.ETW_DNS['provider_guid']
 REQUEST_RESP_EVENT_ID: int = const.ETW_DNS['event_ids']['dns_request_response']
 
-WAIT_FOR_PROCESS_POLLER_PID_SECONDS: int = 3
-WAIT_FOR_PROCESS_POLLER_PID_COUNTS: int = WAIT_FOR_PROCESS_POLLER_PID_SECONDS * 10
-
 
 class DnsRequestResponseTrace:
     """DnsTrace class use to trace DNS events from Windows Event Tracing for EventId 3008."""
@@ -20,7 +19,8 @@ class DnsRequestResponseTrace:
             attrs: list = None,
             session_name: str = None,
             close_existing_session_name: bool = True,
-            skip_record_list: list = None
+            skip_record_list: list = None,
+            process_pool_shared_dict_proxy: multiprocessing.managers.DictProxy = None
     ):
         """
         :param attrs: List of attributes to return. If None, all attributes will be returned.
@@ -32,6 +32,13 @@ class DnsRequestResponseTrace:
                 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']
+        :param process_pool_shared_dict_proxy: multiprocessing.managers.DictProxy, multiprocessing shared dict proxy
+            that contains current processes.
+            Check the 'atomicshop\process_poller\simple_process_pool.py' SimpleProcessPool class for more information.
+
+            For this specific class it means that you can run the process poller outside of this class and pass the
+            'process_pool_shared_dict_proxy' to this class. Then you can get the process name and command line for
+            the DNS events from the 'process_pool_shared_dict_proxy' and use it also in other classes.
 
         -------------------------------------------------
 
@@ -40,7 +47,7 @@ class DnsRequestResponseTrace:
 
 
             dns_trace_w = dns_trace.DnsTrace(
-                attrs=['pid', 'name', 'cmdline', 'domain', 'query_type'],
+                attrs=['pid', 'name', 'cmdline', 'query', 'query_type'],
                 session_name='MyDnsTrace',
                 close_existing_session_name=True,
                 enable_process_poller=True
@@ -53,6 +60,7 @@ class DnsRequestResponseTrace:
         """
 
         self.attrs = attrs
+        self.process_pool_shared_dict_proxy: multiprocessing.managers.DictProxy = process_pool_shared_dict_proxy
 
         if skip_record_list:
             self.skip_record_list: list = skip_record_list
@@ -68,7 +76,8 @@ class DnsRequestResponseTrace:
             event_id_filters=[REQUEST_RESP_EVENT_ID],
             session_name=session_name,
             close_existing_session_name=close_existing_session_name,
-            enable_process_poller=True
+            enable_process_poller=True,
+            process_pool_shared_dict_proxy=self.process_pool_shared_dict_proxy
         )
 
     def start(self):
@@ -90,13 +99,52 @@ class DnsRequestResponseTrace:
         :return: Dictionary with the event data.
         """
 
+        # Get the event from ETW as is.
         event = self.event_trace.emit()
 
+        # Get the raw query results string from the event.
+        query_results: str = event['EventHeader']['QueryResults']
+
+        if query_results != '':
+            query_results_list: list = query_results.split(';')
+
+            addresses_ips: list = list()
+            addresses_cnames: list = list()
+            for query_result in query_results_list:
+                # If there is a type in the query result, it means it is a cname (domain).
+                if 'type' in query_result:
+                    query_result = query_result.split(' ')[-1]
+
+                # But we'll still make sure that the query result is an IP address or not.
+                if ip_addresses.is_ip_address(query_result):
+                    addresses_ips.append(query_result)
+                # If it is not empty, then it is a cname.
+                elif query_result != '':
+                    addresses_cnames.append(query_result)
+        # if the query results are empty, then we'll just set the addresses to empty lists.
+        else:
+            addresses_ips: list = list()
+            addresses_cnames: list = list()
+
+        status_id: str = str(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 status_id == '0':
+            status = 'Success'
+        else:
+            status = 'Error'
+
         event_dict: dict = {
+            'timestamp': event['timestamp'],
             'event_id': event['EventId'],
-            'domain': event['EventHeader']['QueryName'],
+            'query': event['EventHeader']['QueryName'],
             'query_type_id': str(event['EventHeader']['QueryType']),
             'query_type': dns.TYPES_DICT[str(event['EventHeader']['QueryType'])],
+            'result_ips': ','.join(addresses_ips),
+            'result_cnames': ','.join(addresses_cnames),
+            'status_id': status_id,
+            'status': status,
             'pid': event['pid'],
             'name': event['name'],
             'cmdline': event['cmdline']
@@ -107,42 +155,6 @@ class DnsRequestResponseTrace:
         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['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['EventHeader']['QueryResults'].split(';')
-
-            # Iterating through all the parameters that we got from 'QueryResults' key.
-            for parameter in list_of_parameters:
-                # If 'type' string is present it means that entry is a domain;
-                if 'type' in parameter:
-                    # Remove the 'type' string and get the domain name.
-                    current_iteration_parameter = parameter.rsplit(' ', maxsplit=1)[1]
-                    # Add the variable to the list of other answers.
-                    list_of_other_domains.append(current_iteration_parameter)
-                # If 'type' string is not present it means that entry is an IP.
-                else:
-                    # Sometimes the last parameter in the 'QueryResults' key after ';' character will be empty, skip it.
-                    if parameter:
-                        list_of_ips.append(parameter)
-
-        event_dict['ips'] = list_of_ips
-        event_dict['other_domains'] = list_of_other_domains
-
-        # Getting the 'QueryStatus' key.
-        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['EventHeader']['QueryStatus'] == '0':
-            event_dict['status'] = 'Success'
-        else:
-            event_dict['status'] = 'Error'
-
         if self.attrs:
             event_dict = dicts.reorder_keys(
                 event_dict, self.attrs, skip_keys_not_in_list=True)

atomicshop/etws/traces/trace_tcp.py

@@ -0,0 +1,130 @@
+import multiprocessing.managers
+
+from .. import trace, providers
+from ...basics import dicts
+
+
+ETW_DEFAULT_SESSION_NAME: str = 'AtomicShopTcpTrace'
+
+PROVIDER_NAME: str = "Microsoft-Windows-TCPIP"
+PROVIDER_GUID: str = '{' + providers.get_provider_guid_by_name(PROVIDER_NAME) + '}'
+REQUEST_RESP_EVENT_ID: int = 1033
+
+
+class TcpIpNewConnectionsTrace:
+    """
+    TcpIpNewConnectionsTrace class use to trace new connection events from Windows Event Tracing:
+    Provider: Microsoft-Windows-TCPIP
+    EventId: 1033
+    """
+    def __init__(
+            self,
+            attrs: list = None,
+            session_name: str = None,
+            close_existing_session_name: bool = True,
+            process_pool_shared_dict_proxy: multiprocessing.managers.DictProxy = None
+    ):
+        """
+        :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.
+            True: if ETW session with 'session_name' exists, you will be notified and the session will be closed.
+                Then the new session with this name will be created.
+            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 process_pool_shared_dict_proxy: multiprocessing.managers.DictProxy, multiprocessing shared dict proxy
+            that contains current processes.
+            Check the 'atomicshop\process_poller\simple_process_pool.py' SimpleProcessPool class for more information.
+
+            For this specific class it means that you can run the process poller outside of this class and pass the
+            'process_pool_shared_dict_proxy' to this class. Then you can get the process name and command line for
+            the DNS events from the 'process_pool_shared_dict_proxy' and use it also in other classes.
+
+        -------------------------------------------------
+
+        Usage Example:
+            from atomicshop.etw import tcp_trace
+
+
+            tcp_trace_w = tcp_trace.TcpIpNewConnectionsTrace(
+                attrs=['pid', 'name', 'cmdline', 'domain', 'query_type'],
+                session_name='MyTcpTrace',
+                close_existing_session_name=True,
+                enable_process_poller=True
+            )
+            tcp_trace_w.start()
+            while True:
+                tcp_dict = tcp_trace_w.emit()
+                print(tcp_dict)
+            tcp_trace_w.stop()
+        """
+
+        self.attrs = attrs
+        self.process_pool_shared_dict_proxy: multiprocessing.managers.DictProxy = process_pool_shared_dict_proxy
+
+        if not session_name:
+            session_name = ETW_DEFAULT_SESSION_NAME
+
+        self.event_trace = trace.EventTrace(
+            providers=[(PROVIDER_NAME, PROVIDER_GUID)],
+            # 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,
+            enable_process_poller=True,
+            process_pool_shared_dict_proxy=self.process_pool_shared_dict_proxy
+        )
+
+    def start(self):
+        self.event_trace.start()
+
+    def stop(self):
+        self.event_trace.stop()
+
+    def emit(self):
+        """
+        Function that will return the next event from the queue.
+        The queue is blocking, so if there is no event in the queue, the function will wait until there is one.
+
+        Usage Example:
+            while True:
+                tcp_dict = tcp_trace.emit()
+                print(tcp_dict)
+
+        :return: Dictionary with the event data.
+        """
+
+        # Get the event from ETW as is.
+        event = self.event_trace.emit()
+
+        local_address_port: str = event['EventHeader']['LocalAddress']
+        remote_address_port: str = event['EventHeader']['RemoteAddress']
+
+        if 'ffff' in local_address_port:
+            pass
+
+        local_address, local_port = local_address_port.rsplit(':', 1)
+        local_address = local_address.replace('[', '').replace(']', '')
+
+        remote_address, remote_port = remote_address_port.rsplit(':', 1)
+        remote_address = remote_address.replace('[', '').replace(']', '')
+
+        event_dict: dict = {
+            'timestamp': event['timestamp'],
+            'event_id': event['EventId'],
+            'local_ip': local_address,
+            'local_port': local_port,
+            'remote_ip': remote_address,
+            'remote_port': remote_port,
+            'status': event['EventHeader']['Status'],
+            'pid': event['pid'],
+            'name': event['name'],
+            'cmdline': event['cmdline']
+        }
+
+        if self.attrs:
+            event_dict = dicts.reorder_keys(
+                event_dict, self.attrs, skip_keys_not_in_list=True)
+
+        return event_dict

atomicshop/file_io/csvs.py

@@ -2,11 +2,10 @@ import csv
 import io
 from typing import Tuple, List
 
-from .file_io import read_file_decorator
 from . import file_io
 
 
-@read_file_decorator
+@file_io.read_file_decorator
 def read_csv_to_list_of_dicts_by_header(
         file_path: str,
         file_mode: str = 'r',
@@ -53,7 +52,7 @@ def read_csv_to_list_of_dicts_by_header(
     return csv_list, header
 
 
-@read_file_decorator
+@file_io.read_file_decorator
 def read_csv_to_list_of_lists(
         file_path: str,
         file_mode: str = 'r',
@@ -103,7 +102,8 @@ def read_csv_to_list_of_lists(
 def write_list_to_csv(
         content_list: list,
         file_path: str,
-        mode: str = 'w'
+        mode: str = 'w',
+        encoding: str = None
 ) -> None:
     """
     This function got dual purpose:
@@ -115,10 +115,12 @@ def write_list_to_csv(
     :param content_list: List object that each iteration contains dictionary with same keys and different values.
     :param file_path: Full file path to CSV file.
     :param mode: String, file writing mode. Default is 'w'.
+    :param encoding: String, encoding of the file. Default is 'None'.
+        Example: 'utf-8', 'utf-16', 'cp1252'.
     :return: None.
     """
 
-    with open(file_path, mode=mode, newline='') as csv_file:
+    with open(file_path, mode=mode, newline='', encoding=encoding) as csv_file:
         if len(content_list) > 0 and isinstance(content_list[0], dict):
             # Treat the list as list of dictionaries.
             header = content_list[0].keys()
@@ -254,3 +256,23 @@ def escape_csv_line_to_list(csv_line: list) -> list:
         result_csv_entries.append(escape_csv_value(entry))
 
     return result_csv_entries
+
+
+def get_number_of_cells_in_string_line(line: str) -> int:
+    """
+    Function to get number of cells in CSV line.
+
+    :param line: String, line of CSV file.
+    :return: int, number of cells in the line.
+    """
+
+    # Create CSV reader from 'input_file'. By default, the first row will be the header if 'fieldnames' is None.
+    csv_reader = csv.reader([line])
+
+    # Get the first row of the CSV file.
+    csv_list = list(csv_reader)
+
+    # Get the number of cells in the first row.
+    number_of_cells = len(csv_list[0])
+
+    return number_of_cells

atomicshop/file_io/docxs.py

@@ -12,7 +12,7 @@ def get_hyperlinks(docx_path):
     :return: list of strings, hyperlinks.
     """
 
-    hyperlinks: list = list()
+    hyperlinks: list[dict] = list()
 
     try:
         doc = Document(docx_path)
@@ -26,7 +26,19 @@ def get_hyperlinks(docx_path):
         if not paragraph.hyperlinks:
             continue
         for hyperlink in paragraph.hyperlinks:
-            hyperlinks.append(hyperlink.address)
+            # Hyperlinks are stored in docx (document.xml) without the fragment part.
+            # Fragment is the anchor of the link, for example: 'https://www.example.com#anchor'.
+            # So the hyperlink.address is stored as 'https://www.example.com'.
+            # And the fragment is stored in the hyperlink.fragment as 'anchor'.
+            # For the full hyperlink, we need to concatenate the address and the fragment.
+            # If there is no anchor in the link the fragment will be empty string ('').
+            # Basically, we don't need to add the fragment to the hyperlink if it's empty, we can just use the url.
+            # if hyperlink.fragment:
+            #     hyperlinks.append(hyperlink.address + "#" + hyperlink.fragment)
+            hyperlinks.append({
+                'url': hyperlink.url,
+                'text': hyperlink.text
+            })
 
     return hyperlinks
 
@@ -66,24 +78,27 @@ def search_for_hyperlink_in_files(directory_path: str, hyperlink: str, relative_
         raise NotADirectoryError(f"Directory doesn't exist: {directory_path}")
 
     # Get all the docx files in the specified directory.
-    files = filesystem.get_file_paths_from_directory(
-        directory_path, file_name_check_pattern="*\.docx",
+    files = filesystem.get_paths_from_directory(
+        directory_path, get_file=True, file_name_check_pattern="*.docx",
         add_relative_directory=True, relative_file_name_as_directory=True)
 
-    found_in_files: list = list()
+    found_in_files: list[dict] = list()
     for file_path in files:
-        hyperlinks = get_hyperlinks(file_path['file_path'])
-        if hyperlink in hyperlinks:
-            found_in_files.append(file_path)
+        doc_hyperlinks = get_hyperlinks(file_path.path)
+        for doc_link in doc_hyperlinks:
+            if hyperlink in doc_link['url']:
+                if relative_paths:
+                    path: str = file_path.relative_dir
+                else:
+                    path: str = file_path.path
 
-    if relative_paths:
-        result_list = list()
-        for found_file in found_in_files:
-            result_list.append(found_file['relative_dir'])
-    else:
-        result_list = found_in_files
+                found_in_files.append({
+                    'path':path,
+                    'link':doc_link['url'],
+                    'text':doc_link['text']
+                })
 
-    return result_list
+    return found_in_files
 
 
 def search_for_hyperlink_in_files_interface_main(script_directory: str = None):
@@ -126,10 +141,12 @@ def search_for_hyperlink_in_files_interface_main(script_directory: str = None):
     found_in_files = search_for_hyperlink_in_files(
         config['directory_path'], config['hyperlink'], relative_paths=config['relative_paths'])
 
-    print_api(f"Found in [{len(found_in_files)}] files:", color="blue")
+    print_api(f"Found [{len(found_in_files)}] links:", color="blue")
 
     for index, found_file in enumerate(found_in_files):
         print_api(f"[{index+1}]", print_end="", color="green")
-        print_api(f" {found_file}")
+        print_api(f" {found_file['path']}")
+        print_api(f" {found_file['link']}", color="cyan")
+        print_api(f" {found_file['text']}", color="orange")
 
     input('[*] Press [Enter] to exit...')