monkeytoolbox 0.1.0__py3-none-any.whl

CHANGELOG.md

@@ -0,0 +1,15 @@
+# Changelog
+All notable changes to this project will be documented in this
+file.
+
+The format is based on [Keep a
+Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to
+the [PEP 440 version scheme](https://peps.python.org/pep-0440/#version-scheme).
+
+
+## [Unreleased]
+### Added
+### Changed
+### Fixed
+### Removed
+### Security

README.md

@@ -0,0 +1,14 @@
+# monkeytoolbox
+
+This project contains a collection of convenience functions, classes, and
+utilities written in Python. It is used mainly by [Infection
+Monkey](https://github.com/guardicore/monkey).
+
+## Installation
+`pip install monkeytoolbox`
+
+## Running tests
+```
+$> poetry install
+$> poetry run pytest
+```

monkeytoolbox/__init__.py

@@ -0,0 +1,34 @@
+from .environment import get_os, get_hardware_id
+from .decorators import request_cache
+from .code_utils import (
+    apply_filters,
+    queue_to_list,
+    del_key,
+    insecure_generate_random_string,
+    secure_generate_random_string,
+    PeriodicCaller
+)
+from .file_utils import (
+    append_bytes,
+    make_fileobj_copy,
+    InvalidPath,
+    expand_path,
+    get_text_file_contents,
+    get_all_regular_files_in_directory,
+    get_binary_io_sha256_hash
+)
+from .secure_directory import create_secure_directory
+from .secure_file import open_new_securely_permissioned_file
+from .threading import (
+    ThreadSafeIterator, 
+    InterruptableThreadMixin, 
+    interruptible_function, 
+    interruptible_iter,
+    create_daemon_thread,
+    run_worker_threads
+)
+from .network_utils import (
+    port_is_used, 
+    get_network_interfaces, 
+    get_my_ip_addresses
+)

monkeytoolbox/code_utils.py

@@ -0,0 +1,160 @@
+import logging
+import queue
+import random
+import secrets
+import string
+from threading import Event, Thread
+from typing import Any, Callable, Iterable, List, MutableMapping, Optional, TypeVar
+
+T = TypeVar("T")
+
+logger = logging.getLogger(__name__)
+
+
+def apply_filters(filters: Iterable[Callable[[T], bool]], iterable: Iterable[T]) -> Iterable[T]:
+    """
+    Applies multiple filters to an iterable
+
+    :param filters: An iterable of filters to be applied to the iterable
+    :param iterable: An iterable to be filtered
+    :return: A new iterable with the filters applied
+    """
+    filtered_iterable = iterable
+    for f in filters:
+        filtered_iterable = filter(f, filtered_iterable)
+
+    return filtered_iterable
+
+
+def queue_to_list(q: queue.Queue) -> List[Any]:
+    list_ = []
+    try:
+        while True:
+            list_.append(q.get_nowait())
+    except queue.Empty:
+        pass
+
+    return list_
+
+
+def del_key(mapping: MutableMapping[T, Any], key: T):
+    """
+    Delete a key from a mapping.
+
+    Unlike the `del` keyword, this function does not raise a KeyError
+    if the key does not exist.
+
+    :param mapping: A mapping from which a key will be deleted
+    :param key: A key to delete from `mapping`
+    """
+    mapping.pop(key, None)
+
+
+def insecure_generate_random_string(
+    n: int, character_set: str = string.ascii_letters + string.digits
+) -> str:
+    """
+    Generate a random string
+
+    This function generates a random string. The length is specified by the user. The character set
+    can optionally be specified by the user.
+
+    WARNING: This function is not safe to use for cryptographic purposes.
+
+    :param n: The desired number of characters in the random string
+    :param character set: The set of characters that may be included in the random string, defaults
+                          to alphanumerics
+    """
+    return _generate_random_string(random.choices, n, character_set)  # noqa: DUO102
+
+
+def secure_generate_random_string(
+    n: int, character_set: str = string.ascii_letters + string.digits
+) -> str:
+    """
+    Generate a random string
+
+    This function generates a random string. The length is specified by the user. The character set
+    can optionally be specified by the user.
+
+    This function is safe to use for cryptographic purposes.
+
+    WARNING: This function may block if the system does not have sufficient entropy.
+
+    :param n: The desired number of characters in the random string
+    :param character set: The set of characters that may be included in the random string, defaults
+                          to alphanumerics
+    """
+    return _generate_random_string(secrets.SystemRandom().choices, n, character_set)
+
+
+# Note: Trying to typehint the rng parameter is more trouble than it's worth
+def _generate_random_string(rng, n: int, character_set: str) -> str:
+    return "".join(rng(character_set, k=n))
+
+
+class PeriodicCaller:
+    """
+    Periodically calls a function
+
+    Given a callable and a period, this component calls the callback periodically. The calls can
+    occur in the background by calling the `start()` method, or in the foreground by calling the
+    `run()` method. Note that this component is susceptible to "timer creep". In other words, the
+    callable is not called every `period` seconds. It is called `period` seconds after the last call
+    completes. This prevents multiple calls to the callback occurring concurrently.
+    """
+
+    def __init__(self, callback: Callable[[], None], period: float, name: Optional[str] = None):
+        """
+        :param callback: A callable to be called periodically
+        :param period: The time to wait between calls of `callback`.
+        :param name: A human-readable name for this caller that will be used in debug logging
+        """
+        self._callback = callback
+        self._period = period
+
+        self._name = f"PeriodicCaller-{callback.__name__}" if name is None else name
+
+        self._stop = Event()
+        self._thread: Optional[Thread] = None
+
+    def start(self):
+        """
+        Periodically call the callback in the background
+        """
+        logger.debug(f"Starting {self._name}")
+
+        self._stop.clear()
+        self._thread = Thread(daemon=True, name=self._name, target=self.run)
+        self._thread.start()
+
+    def run(self):
+        """
+        Periodically call the callback and block until `stop()` is called
+        """
+        logger.debug(f"Successfully started {self._name}")
+
+        while not self._stop.is_set():
+            self._callback()
+            self._stop.wait(self._period)
+
+        logger.debug(f"Successfully stopped {self._name}")
+
+    def stop(self, timeout: Optional[float] = None):
+        """
+        Stop this component from making any further calls
+
+        When the timeout argument is not present or None, the operation will block until the
+        PeriodicCaller stops.
+
+        :param timeout: The number of seconds to wait for this component to stop
+        """
+        logger.debug(f"Stopping {self._name}")
+
+        self._stop.set()
+
+        if self._thread is not None:
+            self._thread.join(timeout=timeout)
+
+            if self._thread.is_alive():
+                logger.warning(f"Timed out waiting for {self._name} to stop")

monkeytoolbox/decorators.py

@@ -0,0 +1,66 @@
+import threading
+from functools import wraps
+from typing import Any, Callable
+
+from egg_timer import EggTimer
+
+
+def request_cache(ttl: float):
+    """
+    This is a decorator that allows a single response of a function to be cached with an expiration
+    time (TTL). The first call to the function is executed and the response is cached. Subsequent
+    calls to the function result in the cached value being returned until the TTL elapses. Once the
+    TTL elapses, the cache is considered stale and the decorated function will be called, its
+    response cached, and the TTL reset.
+
+    An example usage of this decorator is to wrap a function that makes frequent slow calls to an
+    external resource, such as an HTTP request to a remote endpoint. If the most up-to-date
+    information is not need, this decorator provides a simple way to cache the response for a
+    certain amount of time.
+
+    Example:
+        @request_cache(600)
+        def raining_outside():
+            return requests.get(f"https://weather.service.api/check_for_rain/{MY_ZIP_CODE}")
+
+    The request cache can be manually cleared if desired:
+        status_1 = raining_outside()
+        status_2 = raining_outside()
+        raining_outside.clear_cache()
+        status_3 = raining_outside()
+
+        assert status_1 == status_2
+        assert status_1 != status_3
+
+    :param ttl: The time-to-live in seconds for the cached return value
+    :return: The return value of the decorated function, or the cached return value if the TTL has
+             not elapsed.
+    """
+
+    def decorator(fn: Callable) -> Callable:
+        cached_value = None
+        timer = EggTimer()
+        lock = threading.Lock()
+
+        @wraps(fn)
+        def wrapper(*args, **kwargs) -> Any:
+            nonlocal cached_value, timer, lock
+
+            with lock:
+                if timer.is_expired():
+                    cached_value = fn(*args, **kwargs)
+                    timer.set(ttl)
+
+            return cached_value
+
+        def clear_cache():
+            nonlocal timer, lock
+
+            with lock:
+                timer.set(0)
+
+        wrapper.clear_cache = clear_cache  # type: ignore [attr-defined]
+
+        return wrapper
+
+    return decorator

monkeytoolbox/environment.py

@@ -0,0 +1,38 @@
+import platform
+import uuid
+from contextlib import suppress
+
+from monkeytypes import HardwareID, OperatingSystem
+
+
+def get_os() -> OperatingSystem:
+    if platform.system() == "Windows":
+        return OperatingSystem.WINDOWS
+
+    return OperatingSystem.LINUX
+
+
+def get_hardware_id() -> HardwareID:
+    if get_os() == OperatingSystem.WINDOWS:
+        return _get_hardware_id_windows()
+
+    return _get_hardware_id_linux()
+
+
+def _get_hardware_id_windows() -> HardwareID:
+    return uuid.getnode()
+
+
+def _get_hardware_id_linux() -> HardwareID:
+    # Different compile-time parameters for Python on Linux can cause `uuid.getnode()` to yield
+    # different results. Calling `uuid._ip_getnode()` directly seems to be the most reliable way to
+    # get consistend IDs across different Python binaries. See
+    # https://github.com/guardicore/monkey/issues/3176 for more details
+
+    with suppress(AttributeError):
+        machine_id = uuid._ip_getnode()  # type: ignore [attr-defined]
+
+    if machine_id is None:
+        machine_id = uuid.getnode()
+
+    return machine_id

monkeytoolbox/file_utils.py

@@ -0,0 +1,78 @@
+import hashlib
+import io
+import logging
+import os
+import shutil
+from pathlib import Path
+from typing import BinaryIO, Iterable
+
+logger = logging.getLogger(__name__)
+
+
+MAX_BLOCK_SIZE = 65536
+
+
+class InvalidPath(Exception):
+    pass
+
+
+def expand_path(path: str) -> Path:
+    if not path:
+        raise InvalidPath("Empty path provided")
+
+    return Path(os.path.expandvars(os.path.expanduser(path)))
+
+
+def get_binary_io_sha256_hash(binary: BinaryIO) -> str:
+    """
+    Calculates sha256 hash from a file-like object
+
+    :param binary: file-like object from which we calculate the hash
+    :return: sha256 hash from the file-like object
+    """
+    sha256 = hashlib.sha256()
+    for block in iter(lambda: binary.read(MAX_BLOCK_SIZE), b""):
+        sha256.update(block)
+
+    return sha256.hexdigest()
+
+
+def get_all_regular_files_in_directory(dir_path: Path) -> Iterable[Path]:
+    return filter(lambda f: f.is_file(), dir_path.iterdir())
+
+
+def get_text_file_contents(file_path: Path) -> str:
+    with open(file_path, "rt") as f:
+        file_contents = f.read()
+    return file_contents
+
+
+def make_fileobj_copy(src: BinaryIO) -> BinaryIO:
+    """
+    Creates a file-like object that is a copy of the provided file-like object
+
+    The source file-like object is reset to position 0 and a copy is made. Both the source file and
+    the copy are reset to position 0 before returning.
+
+    :param src: A file-like object to copy
+    :return: A file-like object that is a copy of the provided file-like object
+    """
+    dst = io.BytesIO()
+
+    src.seek(0)
+    shutil.copyfileobj(src, dst)
+
+    src.seek(0)
+    dst.seek(0)
+
+    return dst
+
+
+def append_bytes(file: BinaryIO, bytes_to_append: bytes) -> BinaryIO:
+    starting_position = file.tell()
+
+    file.seek(0, io.SEEK_END)
+    file.write(bytes_to_append)
+    file.seek(starting_position, io.SEEK_SET)
+
+    return file

monkeytoolbox/network_utils.py

@@ -0,0 +1,53 @@
+import ipaddress
+from ipaddress import IPv4Address, IPv4Interface
+from typing import Iterable, List, Optional, Sequence
+
+import ifaddr
+import psutil
+
+
+def get_my_ip_addresses() -> Sequence[IPv4Address]:
+    return [interface.ip for interface in get_network_interfaces()]
+
+
+def get_network_interfaces() -> List[IPv4Interface]:
+    local_interfaces = []
+    for adapter in ifaddr.get_adapters():
+        for ip in _select_ipv4_ips(adapter.ips):
+            interface = ipaddress.IPv4Interface(f"{ip.ip}/{ip.network_prefix}")
+            if not interface.ip.is_link_local:
+                local_interfaces.append(interface)
+
+    return local_interfaces
+
+
+def _select_ipv4_ips(ips: Iterable[ifaddr.IP]) -> Iterable[ifaddr.IP]:
+    return filter(lambda ip: _is_ipv4(ip) and ip.ip != "127.0.0.1", ips)
+
+
+def _is_ipv4(ip: ifaddr.IP) -> bool:
+    # In ifaddr, IPv4 addresses are strings, while IPv6 addresses are tuples
+    return type(ip.ip) is str
+
+
+def port_is_used(
+    port: int,
+    ip_addresses: Optional[Sequence[IPv4Address]],
+) -> bool:
+    connections = get_connections([port], ip_addresses)
+    return len(connections) > 0
+
+
+def get_connections(
+    ports: Optional[Sequence[int]] = None,
+    ip_addresses: Optional[Sequence[IPv4Address]] = None,
+) -> List[psutil._common.sconn]:
+    connections = psutil.net_connections()
+    if ports:
+        connections = [connection for connection in connections if connection.laddr.port in ports]
+    if ip_addresses:
+        ip_addresses_ = list(map(str, ip_addresses))
+        connections = [
+            connection for connection in connections if connection.laddr.ip in ip_addresses_
+        ]
+    return connections

monkeytoolbox/secure_directory.py

@@ -0,0 +1,88 @@
+import logging
+import stat
+from pathlib import Path
+from typing import Callable
+
+from . import get_os
+
+from monkeytypes import OperatingSystem
+
+if get_os() == OperatingSystem.WINDOWS:
+    import win32file
+    import win32security
+
+    from .windows_permissions import get_security_descriptor_for_owner_only_permissions
+
+logger = logging.getLogger(__name__)
+
+
+class FailedDirectoryCreationError(Exception):
+    pass
+
+
+def create_secure_directory(path: Path):
+    if get_os() == OperatingSystem.WINDOWS:
+        make_existing_directory_secure_for_os = _make_existing_directory_secure_windows
+        create_secure_directory_for_os = _create_secure_directory_windows
+    else:
+        make_existing_directory_secure_for_os = _make_existing_directory_secure_linux
+        create_secure_directory_for_os = _create_secure_directory_linux
+
+    if path.exists():
+        _check_path_is_directory(path)
+        _make_existing_directory_secure(make_existing_directory_secure_for_os, path)
+    else:
+        _create_secure_directory(create_secure_directory_for_os, path)
+
+
+def _check_path_is_directory(path: Path):
+    if not path.is_dir():
+        raise FailedDirectoryCreationError(
+            f'The path "{path}" already exists and is not a directory'
+        )
+
+    logger.info(f"A directory already exists at {path}")
+
+
+def _make_existing_directory_secure(fn_for_os: Callable, path: Path):
+    try:
+        fn_for_os(path)
+    except Exception as err:
+        message = (
+            "An error occured while changing the existing directory's permissions"
+            f"to be secure: {str(err)}"
+        )
+        logger.exception(message)
+        raise FailedDirectoryCreationError(err)
+
+
+def _create_secure_directory(fn_for_os: Callable, path: Path):
+    try:
+        fn_for_os(path)
+    except Exception as err:
+        message = f"Could not create a secure directory at {path}: {str(err)}"
+        logger.error(message)
+        raise FailedDirectoryCreationError(message)
+
+
+def _make_existing_directory_secure_windows(path: Path):
+    security_descriptor = get_security_descriptor_for_owner_only_permissions()
+    win32security.SetFileSecurity(
+        str(path), win32security.DACL_SECURITY_INFORMATION, security_descriptor
+    )
+
+
+def _create_secure_directory_windows(path: Path):
+    security_attributes = win32security.SECURITY_ATTRIBUTES()
+    security_attributes.SECURITY_DESCRIPTOR = (
+        get_security_descriptor_for_owner_only_permissions()
+    )
+    win32file.CreateDirectory(str(path), security_attributes)
+
+
+def _make_existing_directory_secure_linux(path: Path):
+    path.chmod(mode=stat.S_IRWXU)
+
+
+def _create_secure_directory_linux(path: Path):
+    path.mkdir(mode=stat.S_IRWXU)

monkeytoolbox/secure_file.py

@@ -0,0 +1,85 @@
+import logging
+import os
+import stat
+from contextlib import contextmanager
+from typing import Generator
+
+from . import get_os
+from monkeytypes import OperatingSystem
+
+if get_os() == OperatingSystem.WINDOWS:
+    import win32file
+    import win32job
+    import win32security
+
+    from .windows_permissions import get_security_descriptor_for_owner_only_permissions
+
+
+logger = logging.getLogger(__name__)
+
+
+@contextmanager
+def open_new_securely_permissioned_file(path: str, mode: str = "w") -> Generator:
+    if get_os() == OperatingSystem.WINDOWS:
+        # TODO: Switch from string to Path object to avoid this hack.
+        fd = _get_file_descriptor_for_new_secure_file_windows(str(path))
+    else:
+        fd = _get_file_descriptor_for_new_secure_file_linux(path)
+
+    with open(fd, mode) as f:
+        yield f
+
+
+def _get_file_descriptor_for_new_secure_file_linux(path: str) -> int:
+    try:
+        mode = stat.S_IRUSR | stat.S_IWUSR
+        flags = (
+            os.O_RDWR | os.O_CREAT | os.O_EXCL
+        )  # read/write, create new, throw error if file exists
+        fd = os.open(path, flags, mode)
+
+        return fd
+
+    except Exception as ex:
+        logger.error(f'Could not create a file at "{path}": {str(ex)}')
+        raise ex
+
+
+def _get_file_descriptor_for_new_secure_file_windows(path: str) -> int:
+    try:
+        file_access = win32file.GENERIC_READ | win32file.GENERIC_WRITE
+
+        # Enables other processes to open this file with read-only access.
+        # Attempts by other processes to open the file for writing while this
+        # process still holds it open will fail.
+        file_sharing = win32file.FILE_SHARE_READ
+
+        security_attributes = win32security.SECURITY_ATTRIBUTES()
+        security_attributes.SECURITY_DESCRIPTOR = (
+            get_security_descriptor_for_owner_only_permissions()
+        )
+        file_creation = win32file.CREATE_NEW  # fails if file exists
+        file_attributes = win32file.FILE_FLAG_BACKUP_SEMANTICS
+
+        handle = win32file.CreateFile(
+            path,
+            file_access,
+            file_sharing,
+            security_attributes,
+            file_creation,
+            file_attributes,
+            _get_null_value_for_win32(),
+        )
+
+        detached_handle = handle.Detach()
+
+        return win32file._open_osfhandle(detached_handle, os.O_RDWR)
+
+    except Exception as ex:
+        logger.error(f'Could not create a file at "{path}": {str(ex)}')
+        raise ex
+
+
+def _get_null_value_for_win32():
+    # https://stackoverflow.com/questions/46800142/in-python-with-pywin32-win32job-the-createjobobject-function-how-do-i-pass-nu  # noqa: E501
+    return win32job.CreateJobObject(None, "")