monkeytoolbox 0.1.0__py3-none-any.whl

monkeytoolbox/threading.py

@@ -0,0 +1,140 @@
+import logging
+import threading
+from functools import wraps
+from itertools import count
+from threading import Lock, Thread
+from typing import Any, Callable, Iterable, Iterator, Optional, Tuple, TypeVar
+
+from monkeytypes import Event
+
+logger = logging.getLogger(__name__)
+
+
+def run_worker_threads(
+    target: Callable[..., None],
+    name_prefix: str,
+    args: Tuple = (),
+    num_workers: int = 2,
+):
+    worker_threads = []
+    counter = run_worker_threads.counters.setdefault(name_prefix, count(start=1))
+    for i in range(0, num_workers):
+        name = f"{name_prefix}-{next(counter):02d}"
+        t = create_daemon_thread(target=target, name=name, args=args)
+        t.start()
+        worker_threads.append(t)
+
+    for t in worker_threads:
+        t.join()
+
+
+run_worker_threads.counters = {}
+
+
+def create_daemon_thread(target: Callable[..., None], name: str, args: Tuple = ()) -> Thread:
+    return Thread(target=target, name=name, args=args, daemon=True)
+
+
+def interruptible_iter(
+    iterator: Iterable,
+    interrupt: Event,
+    log_message: Optional[str] = None,
+    log_level: int = logging.DEBUG,
+) -> Any:
+    """
+    Wraps an iterator so that the iterator can be interrupted if the `interrupt` event is set. This
+    is a convinient way to make loops interruptible and avoids the need to add an `if` to each and
+    every loop.
+    :param Iterable iterator: An iterator that will be made interruptible.
+    :param Event interrupt: An `Event` that, if set, will prevent the remainder of the
+                            iterator's items from being processed.
+    :param str log_message: A message to be logged if the iterator is interrupted. If `log_message`
+                            is `None` (default), then no message is logged.
+    :param int log_level: The log level at which to log `log_message`, defaults to `logging.DEBUG`.
+    """
+    for i in iterator:
+        if interrupt.is_set():
+            if log_message:
+                logger.log(log_level, log_message)
+
+            break
+
+        yield i
+
+
+def interruptible_function(*, msg: Optional[str] = None, default_return_value: Any = None):
+    """
+    This decorator allows a function to be skipped if an interrupt (`Event`) is set. This is
+    useful for interrupting running code without introducing duplicate `if` checks at the beginning
+    of each function.
+
+    Note: It is required that the decorated function accept a keyword-only argument named
+    "interrupt".
+
+    Example:
+        def run_algorithm(*inputs, interrupt: Event):
+            return_value = do_action_1(inputs[1], interrupt=interrupt)
+            return_value = do_action_2(return_value + inputs[2], interrupt=interrupt)
+            return_value = do_action_3(return_value + inputs[3], interrupt=interrupt)
+
+            return return_value
+
+        @interruptible_function(msg="Interrupt detected, skipping action 1", default_return_value=0)
+        def do_action_1(input, *, interrupt: Event):
+            # Process input
+            ...
+
+        @interruptible_function(msg="Interrupt detected, skipping action 2", default_return_value=0)
+        def do_action_2(input, *, interrupt: Event):
+            # Process input
+            ...
+
+        @interruptible_function(msg="Interrupt detected, skipping action 2", default_return_value=0)
+        def do_action_2(input, *, interrupt: Event):
+            # Process input
+            ...
+
+    :param str msg: A message to log at the debug level if an interrupt is detected. Defaults to
+                    None.
+    :param Any default_return_value: A value to return if the wrapped function is not run. Defaults
+                                     to None.
+    """
+
+    def _decorator(fn):
+        @wraps(fn)
+        def _wrapper(*args, interrupt: Event, **kwargs):
+            if interrupt.is_set():
+                if msg:
+                    logger.debug(msg)
+                return default_return_value
+
+            return fn(*args, interrupt=interrupt, **kwargs)
+
+        return _wrapper
+
+    return _decorator
+
+
+class InterruptableThreadMixin:
+    def __init__(self):
+        self._interrupted = threading.Event()
+
+    def stop(self):
+        """Stop a running thread."""
+        self._interrupted.set()
+
+
+T = TypeVar("T")
+
+
+class ThreadSafeIterator(Iterator[T]):
+    """Provides a thread-safe iterator that wraps another iterator"""
+
+    def __init__(self, iterator: Iterator[T]):
+        self._lock = Lock()
+        self._iterator = iterator
+
+    def __next__(self) -> T:
+        while True:
+            with self._lock:
+                return next(self._iterator)

monkeytoolbox/windows_permissions.py

@@ -0,0 +1,52 @@
+from typing import Any, Mapping
+
+import ntsecuritycon
+import win32api
+import win32security
+
+ACCESS_MODE_GRANT_ACCESS = win32security.GRANT_ACCESS
+ACCESS_PERMISSIONS_FULL_CONTROL = ntsecuritycon.FILE_ALL_ACCESS
+INHERITANCE_OBJECT_AND_CONTAINER = (
+    win32security.CONTAINER_INHERIT_ACE | win32security.OBJECT_INHERIT_ACE
+)
+
+
+def get_security_descriptor_for_owner_only_permissions():
+    user_sid = _get_user_pySID_object()
+    entries = [_get_ace_for_owner_only_permissions(user_sid)]
+
+    dacl = win32security.ACL()
+    dacl.SetEntriesInAcl(entries)
+
+    security_descriptor = win32security.SECURITY_DESCRIPTOR()
+    security_descriptor.SetSecurityDescriptorControl(
+        win32security.SE_DACL_PROTECTED, win32security.SE_DACL_PROTECTED
+    )
+    security_descriptor.SetSecurityDescriptorDacl(1, dacl, 0)
+
+    return security_descriptor
+
+
+def _get_user_pySID_object():
+    # Note: We do this using the process handle and token instead of by account name, as some SIDs
+    # have no corresponding account name, such as a logon SID that identifies a logon session. This
+    # is particularly an issue when using the SMB exploiter. See issue #3173.
+    process_handle = win32api.GetCurrentProcess()
+    token_handle = win32security.OpenProcessToken(process_handle, win32security.TOKEN_ALL_ACCESS)
+
+    sid, _ = win32security.GetTokenInformation(token_handle, win32security.TokenUser)
+
+    return sid
+
+
+def _get_ace_for_owner_only_permissions(user_sid) -> Mapping[str, Any]:
+    return {
+        "AccessMode": ACCESS_MODE_GRANT_ACCESS,
+        "AccessPermissions": ACCESS_PERMISSIONS_FULL_CONTROL,
+        "Inheritance": INHERITANCE_OBJECT_AND_CONTAINER,
+        "Trustee": {
+            "TrusteeType": win32security.TRUSTEE_IS_USER,
+            "TrusteeForm": win32security.TRUSTEE_IS_SID,
+            "Identifier": user_sid,
+        },
+    }