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

atomicshop/get_process_list.py

@@ -0,0 +1,133 @@
+from typing import Union, Literal
+
+from .wrappers.pywin32w.wmis import 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 = 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/{process_name_cmd.py → get_process_name_cmd_dll.py}

@@ -1,19 +1,25 @@
-import pkg_resources
+from importlib.resources import files
 import ctypes
 from ctypes import wintypes
 from typing import Literal
+import threading
 
 from .basics import list_of_dicts
 
 
+PACKAGE_DLL_PATH = 'addons/process_list/compiled/Win10x64/process_list.dll'
+FULL_DLL_PATH = str(files(__package__).joinpath(PACKAGE_DLL_PATH))
+
+
 class ProcessNameCmdline:
     def __init__(self, load_dll: bool = False):
         self.load_dll: bool = load_dll
-        self.dll_path: str = pkg_resources.resource_filename(
-            __package__, 'addons/process_list/compiled/Win10x64/process_list.dll')
+        self.dll_path = FULL_DLL_PATH
+
         self.dll = None
         self.callback_output = None
         self.CALLBACKFUNC = None
+        self.CALLBACKFUNC_ref = None
 
         if self.load_dll:
             self.load()
@@ -21,12 +27,16 @@ class ProcessNameCmdline:
     def load(self):
         """
         Load the DLL, initialize the callback function and ctypes.
-        ctypes.WINFUCNTYPE is not thread safe. You should load it inside a thread / process and not outside.
+        ctypes.WINFUNCTYPE is not thread safe. You should load it inside a thread / process and not outside.
         """
         self.dll = ctypes.windll.LoadLibrary(self.dll_path)
         self.callback_output = OutputList()
         self.CALLBACKFUNC = ctypes.WINFUNCTYPE(None, wintypes.DWORD, wintypes.LPWSTR, wintypes.LPWSTR)
+        self.CALLBACKFUNC_ref = self.CALLBACKFUNC(self.callback_output.callback)
+
+        # Set the argument types for the export functions of the DLL.
         self.dll.GetProcessDetails.argtypes = [self.CALLBACKFUNC]
+        self.dll.RequestCancellation.restype = None
 
     def get_process_details(
             self,
@@ -49,23 +59,37 @@ class ProcessNameCmdline:
                 }
         """
 
-        self.dll.GetProcessDetails(self.CALLBACKFUNC(self.callback_output.callback))
+        def enumerate_process():
+            self.dll.GetProcessDetails(self.CALLBACKFUNC_ref)
+
+        thread = threading.Thread(target=enumerate_process)
+        thread.start()
+
+        try:
+            # This is needed to stop the thread if the main thread is interrupted.
+            # If we execute the 'self.dll.GetProcessDetails(self.CALLBACKFUNC_ref)' directly
+            # and we would like to KeyboardInterrupt, we will get an error:
+            # Exception ignored on calling ctypes callback function.
+            thread.join()
 
-        processes = self.callback_output.data
+            processes = self.callback_output.data
 
-        # Clear the callback output list, or it will be appended each time.
-        self.callback_output.data = list()
+            # Clear the callback output list, or it will be appended each time.
+            self.callback_output.data = list()
 
-        if sort_by:
-            processes = list_of_dicts.sort_by_keys(processes, key_list=[sort_by])
+            if sort_by:
+                processes = list_of_dicts.sort_by_keys(processes, key_list=[sort_by])
 
-        if as_dict:
-            processes = convert_processes_to_dict(processes)
+            if as_dict:
+                processes = convert_processes_to_dict(processes)
 
-        return processes
+            return processes
+        except KeyboardInterrupt:
+            self.dll.RequestCancellation()
+            raise
 
 
-def convert_processes_to_dict(process_list: dict) -> dict:
+def convert_processes_to_dict(process_list: list[dict]) -> dict:
     """
     The function will convert the list of processes to dict, while 'pid' values will be converted to key.
     Example:
@@ -99,14 +123,23 @@ class OutputList:
 
     def callback(self, pid, process_name, cmdline):
         try:
+            process_name_decoded = process_name.decode("utf-16") if isinstance(process_name, bytes) else process_name
+            cmdline_decoded = cmdline.decode("utf-16") if isinstance(cmdline, bytes) else (
+                cmdline if cmdline else "Error")
             self.data.append({
                 "pid": pid,
-                "name": process_name.decode("utf-16"),
-                "cmdline": cmdline.decode("utf-16") if cmdline else "Error"
+                "name": process_name_decoded,
+                "cmdline": cmdline_decoded
             })
-        except AttributeError:
+        # except AttributeError:
+        #     self.data.append({
+        #         "pid": pid,
+        #         "name": process_name,
+        #         "cmdline": cmdline if cmdline else "Error"
+        #     })
+        except Exception:
             self.data.append({
                 "pid": pid,
-                "name": process_name,
-                "cmdline": cmdline if cmdline else "Error"
+                "name": "Error",
+                "cmdline": "Error"
             })

atomicshop/http_parse.py

@@ -1,10 +1,43 @@
-# v1.0.1 - 26.03.2023 23:40
 from http.server import BaseHTTPRequestHandler
 from http.client import HTTPResponse
 import http
 from io import BytesIO
 
 
+def get_request_methods() -> list[str]:
+    """
+    Function to return all available HTTP request methods.
+    """
+    # noinspection PyUnresolvedReferences
+    return [method.value for method in http.HTTPMethod]
+
+
+def is_first_bytes_http_request(request_bytes: bytes) -> bool:
+    """
+    Function to check if the first bytes are HTTP request or not.
+    """
+    http_request_methods_list: list = get_request_methods()
+    http_request_methods_list = [method.encode() for method in http_request_methods_list]
+
+    # If the first bytes are HTTP request, then the first word should be one of the HTTP request methods.
+    for method in http_request_methods_list:
+        if request_bytes.startswith(method):
+            return True
+    return False
+
+
+def is_first_bytes_http_response(response_bytes: bytes) -> bool:
+    """
+    Function to check if the first bytes are HTTP response or not.
+    """
+    http_response_beginning: bytes = b'HTTP/'
+
+    # If the first bytes are HTTP response, then the first word should be 'HTTP/'.
+    if response_bytes.startswith(http_response_beginning):
+        return True
+    return False
+
+
 class HTTPRequestParse(BaseHTTPRequestHandler):
     """
     The class will parse HTTP requests.
@@ -49,41 +82,27 @@ class HTTPRequestParse(BaseHTTPRequestHandler):
             client_message.request_raw_decoded = request_decoded
     """
 
-    def __init__(self, request_text):
-        self.rfile = BytesIO(request_text)
-        self.raw_requestline = self.rfile.readline()
-        self.error_code = self.error_message = None
-        self.parse_request()
-
-        # Check if ".path" attribute exists after HTTP request parsing
-        if not hasattr(self, 'path'):
-            self.path = None
+    # noinspection PyMissingConstructor
+    def __init__(self, request_bytes: bytes):
+        self.request_bytes: bytes = request_bytes
 
-        # Before checking for body, we need to make sure that ".headers" property exists, if not, return empty values
-        try:
-            # The "body" of request is in the 'Content-Length' key. If it exists in "headers" - get the body
-            if 'Content-Length' in self.headers.keys():
-                # "self.headers.get('Content-Length')" returns number in string format, "int" converts it to integer
-                self.content_length = int(self.headers.get('Content-Length'))
-                self.body = self.rfile.read(self.content_length)
-            else:
-                self.content_length = None
-                self.body = None
-        except Exception:
-            self.content_length = None
-            self.body = None
-            pass
+        # noinspection PyTypeChecker
+        self.rfile = None
+        self.raw_requestline = None
+        self.error_code = None
+        self.error_message = None
 
-        # Examples:
-        # Getting path: self.path
-        # Getting Request Version: self.request_version
-        # Getting specific header: self.headers['host']
+        self.content_length = None
+        self.body = None
+        # noinspection PyTypeChecker
+        self.path = None
 
+    # noinspection PyMethodOverriding
     def send_error(self, code, message):
         self.error_code = code
         self.error_message = message
 
-    def check_if_http(self):
+    def parse(self):
         """
         Function to check if parsed object is HTTP request or not.
         'reason' will be populated with parsing status and errors.
@@ -117,88 +136,125 @@ class HTTPRequestParse(BaseHTTPRequestHandler):
 
             client_message.request_raw_decoded = request_decoded
         """
-        # Defining variables
-        reason = None
-        function_result: bool = True
-        error: bool = False
-
-        # If there's any error in HTTP parsing
-        if self.error_message:
-            # If error status is "BAD_REQUEST"
-            if self.error_message.startswith("Bad request"):
-                # If it's 'Bad request' this is not HTTP request, so we can
-                # continue the execution and parse the code as NON-HTTP Request.
-                # Currently, seen 'Bad request syntax' and 'Bad request version'.
-                reason = f"HTTP Request Parsing: Not HTTP request: {self.error_message}"
-                function_result = False
-                error: False
-            else:
-                reason = f"HTTP Request Parsing: HTTP Request with Script Undocumented ERROR: {self.error_message}"
-                function_result = True
-                error = True
-        # If there's no error at all in HTTP Parsing, then it's fine HTTP Request
+
+        if self.request_bytes is None or not is_first_bytes_http_request(self.request_bytes):
+            error = "HTTP Request Parsing: Not HTTP request by first bytes."
+            is_http = False
         else:
-            reason = "HTTP Request Parsing: Valid HTTP request"
-            function_result = True
-            error = False
+            error: str = str()
+
+            self.rfile = BytesIO(self.request_bytes)
+            self.raw_requestline = self.rfile.readline()
+            self.error_code = self.error_message = None
+            self.parse_request()
 
-        return function_result, reason, error
+            # Before checking for body, we need to make sure that ".headers" property exists, if not, return empty values.
+            if hasattr(self, 'headers'):
+                # The "body" of request is in the 'Content-Length' key. If it exists in "headers" - get the body
+                if 'Content-Length' in self.headers.keys():
+                    # "self.headers.get('Content-Length')" returns number in string format, "int" converts it to integer
+                    self.content_length = int(self.headers.get('Content-Length'))
+                    self.body = self.rfile.read(self.content_length)
+
+            # Examples:
+            # Getting path: self.path
+            # Getting Request Version: self.request_version
+            # Getting specific header: self.headers['host']
+
+            # If there's any error in HTTP parsing
+            if self.error_message:
+                # If error status is "BAD_REQUEST"
+                if self.error_message.startswith("Bad request"):
+                    # If it's 'Bad request' this is not HTTP request, so we can
+                    # continue the execution and parse the code as NON-HTTP Request.
+                    # Currently, seen 'Bad request syntax' and 'Bad request version'.
+                    error = f"HTTP Request Parsing: Not HTTP request: {self.error_message}"
+                    is_http = False
+                else:
+                    error = f"HTTP Request Parsing: HTTP Request with Script Undocumented ERROR: {self.error_message}"
+                    is_http = False
+            # If there's no error at all in HTTP Parsing, then it's fine HTTP Request
+            else:
+                is_http = True
+
+        return self, is_http, error
 
 
 class FakeSocket:
     """
-    FakeSocket is needed to parse HTTP Response. Socket object is needed for HTTPResponse class input.
+    FakeSocket mimics a socket object for parsing HTTP responses.
     """
     def __init__(self, response_bytes):
         self._file = BytesIO(response_bytes)
 
-    def makefile(self, *args, **kwargs):
+    # noinspection PyUnusedLocal
+    def makefile(self, mode='rb', buffering=-1) -> BytesIO:
+        """
+        Mimics the socket's makefile method, returning the BytesIO object.
+        """
         return self._file
 
+    def fileno(self) -> int:
+        """
+        Provide a dummy file descriptor, as some code might call this.
+        """
+        raise OSError("File descriptor not available in FakeSocket")
+
 
 class HTTPResponseParse:
     def __init__(self, response_raw_bytes: bytes):
-        self.error = None
         self.response_raw_bytes: bytes = response_raw_bytes
-        # Assigning FakeSocket with response_raw_bytes.
-        self.source = FakeSocket(self.response_raw_bytes)
 
-        # Initializing HTTPResponse class with the FakeSocket with response_raw_bytes as input.
-        self.response_raw_decoded = HTTPResponse(self.source)
+        self.error = None
+        self.source = None
+        self.response_raw_parsed = None
+        self.is_http: bool = False
 
-        # Try to parse HTTP Response.
-        try:
-            self.response_raw_decoded.begin()
-        # If there were problems with the status line.
-        except http.client.BadStatusLine:
-            self.error = "HTTP Response Parsing: Not a valid HTTP Response: Bad Status Line."
-            pass
+    def parse(self):
+        if self.response_raw_bytes is None or not is_first_bytes_http_response(self.response_raw_bytes):
+            self.error = "HTTP Response Parsing: Not a valid HTTP Response by first bytes."
+            self.is_http = False
+            self.response_raw_parsed = None
+        else:
+            # Assigning FakeSocket with response_raw_bytes.
+            self.source = FakeSocket(self.response_raw_bytes)
 
-        try:
-            # If no exception was thrown, but there are some problems with headers.
-            if self.response_raw_decoded.headers.defects:
+            # Initializing HTTPResponse class with the FakeSocket with response_raw_bytes as input.
+            # noinspection PyTypeChecker
+            self.response_raw_parsed = HTTPResponse(self.source)
+
+            # Try to parse HTTP Response.
+            try:
+                self.response_raw_parsed.begin()
+                self.is_http = True
+            # If there were problems with the status line.
+            except http.client.BadStatusLine:
+                self.error = "HTTP Response Parsing: Not a valid HTTP Response: Bad Status Line."
+                self.is_http = False
+
+            header_exists: bool = False
+            if (self.response_raw_parsed is not None and hasattr(self.response_raw_parsed, 'headers') and
+                    self.response_raw_parsed.headers is not None):
+                header_exists = True
+
+            if header_exists and self.response_raw_parsed.headers.defects:
                 self.error = f"HTTP Response Parsing: Not a valid HTTP Response: Some defects in headers: " \
-                             f"{self.response_raw_decoded.headers.defects}"
-        # If the attribute of defects doesn't exist, probably the response wasn't parsed at all by the library,
-        # Meaning, that the exception was already handled.
-        except AttributeError:
-            pass
+                             f"{self.response_raw_parsed.headers.defects}"
+                self.is_http = False
 
-        # Before checking for body, we need to make sure that ".headers" property exists, if not, return empty values
-        try:
-            # The "body" of response is in the 'Content-Length' key. If it exists in "headers" - get the body.
-            if 'Content-Length' in self.response_raw_decoded.headers.keys():
-                # "self.response_raw_decoded.headers.get('Content-Length')" returns number in string format,
-                # "int" converts it to integer.
-                self.response_raw_decoded.content_length = int(self.response_raw_decoded.headers.get('Content-Length'))
-                # Basically we need to extract only the number of bytes specified in 'Content-Length' from the end
-                # of the response that we received.
-                # self.response_raw_bytes[-23:]
-                self.response_raw_decoded.body = self.response_raw_bytes[-self.response_raw_decoded.content_length:]
-            else:
-                self.response_raw_decoded.content_length = None
-                self.response_raw_decoded.body = None
-        except Exception:
-            self.response_raw_decoded.content_length = None
-            self.response_raw_decoded.body = None
-            pass
+            if self.is_http:
+                # Before checking for body, we need to make sure that ".headers" property exists,
+                # if not, return empty values.
+                self.response_raw_parsed.content_length = None
+                self.response_raw_parsed.body = None
+                if header_exists and 'Content-Length' in self.response_raw_parsed.headers.keys():
+                    # The "body" of response is in the 'Content-Length' key. If it exists in "headers" - get the body.
+                    # "self.response_raw_decoded.headers.get('Content-Length')" returns number in string format,
+                    # "int" converts it to integer.
+                    self.response_raw_parsed.content_length = int(self.response_raw_parsed.headers.get('Content-Length'))
+                    # Basically we need to extract only the number of bytes specified in 'Content-Length' from the end
+                    # of the response that we received.
+                    # self.response_raw_bytes[-23:]
+                    self.response_raw_parsed.body = self.response_raw_bytes[-self.response_raw_parsed.content_length:]
+
+        return self.response_raw_parsed, self.is_http, self.error

atomicshop/ip_addresses.py

@@ -2,7 +2,12 @@ import ipaddress
 from typing import Union, Literal
 
 
-def is_ip_address(string_value: str, ip_type: Union[Literal['ipv4', 'ipv6'], None] = None) -> bool:
+def is_ip_address(
+        string_value: str,
+        ip_type: Union[
+            Literal['ipv4', 'ipv6'],
+            None] = None
+) -> bool:
     """
     The function checks if the string is an IPv4 or IPv6 address.