maco 1.2.14__py3-none-any.whl → 1.2.16__py3-none-any.whl

model_setup/maco/collector.py

@@ -13,18 +13,20 @@ from multiprocess import Manager, Process, Queue
 from pydantic import BaseModel
 
 from maco import extractor, model, utils, yara
-from maco.exceptions import AnalysisAbortedException
-
-
-class ExtractorLoadError(Exception):
-    pass
-
+from maco.exceptions import AnalysisAbortedException, ExtractorLoadError
 
 logger = logging.getLogger("maco.lib.helpers")
 
 
 def _verify_response(resp: Union[BaseModel, dict]) -> Dict:
-    """Enforce types and verify properties, and remove defaults."""
+    """Enforce types and verify properties, and remove defaults.
+
+    Args:
+        resp (Union[BaseModel, dict])): results from extractor
+
+    Returns:
+        (Dict): results from extractor after verification
+    """
     if not resp:
         return None
     # check the response is valid for its own model
@@ -62,6 +64,8 @@ class ExtractorRegistration(TypedDict):
 
 
 class Collector:
+    """Discover and load extractors from file system."""
+
     def __init__(
         self,
         path_extractors: str,
@@ -70,7 +74,11 @@ class Collector:
         create_venv: bool = False,
         skip_install: bool = False,
     ):
-        """Discover and load extractors from file system."""
+        """Discover and load extractors from file system.
+
+        Raises:
+            ExtractorLoadError: when no extractors are found
+        """
         # maco requires the extractor to be imported directly, so ensure they are available on the path
         full_path_extractors = os.path.abspath(path_extractors)
         full_path_above_extractors = os.path.dirname(full_path_extractors)
@@ -175,7 +183,15 @@ class Collector:
         stream: BinaryIO,
         extractor_name: str,
     ) -> Dict[str, Any]:
-        """Run extractor with stream and verify output matches the model."""
+        """Run extractor with stream and verify output matches the model.
+
+        Args:
+            stream (BinaryIO): Binary stream to analyze
+            extractor_name (str): Name of extractor to analyze stream
+
+        Returns:
+            (Dict[str, Any]): Results from extractor
+        """
         extractor = self.extractors[extractor_name]
         try:
             # Run extractor on a copy of the sample

model_setup/maco/exceptions.py

@@ -1,3 +1,33 @@
+"""Exception classes for extractors."""
+
+
 # Can be raised by extractors to abort analysis of a sample
 # ie. Can abort if preliminary checks at start of run indicate the file shouldn't be analyzed by extractor
-class AnalysisAbortedException(Exception): ...
+class AnalysisAbortedException(Exception):
+    """Raised when extractors voluntarily abort analysis of a sample."""
+
+    pass
+
+
+class ExtractorLoadError(Exception):
+    """Raised when extractors cannot be loaded."""
+
+    pass
+
+
+class InvalidExtractor(ValueError):
+    """Raised when an extractor is invalid."""
+
+    pass
+
+
+class NoHitException(Exception):
+    """Raised when the YARA rule of an extractor doesn't hit."""
+
+    pass
+
+
+class SyntaxError(Exception):
+    """Raised when there's a syntax error in the YARA rule."""
+
+    pass

model_setup/maco/extractor.py

@@ -4,14 +4,8 @@ import logging
 import textwrap
 from typing import BinaryIO, List, Optional, Union
 
-from maco import yara
-
-from . import model
-
-
-class InvalidExtractor(ValueError):
-    pass
-
+from maco import model, yara
+from maco.exceptions import InvalidExtractor
 
 DEFAULT_YARA_RULE = """
 rule {name}
@@ -37,6 +31,11 @@ class Extractor:
     logger: logging.Logger = None  # logger for use when debugging
 
     def __init__(self) -> None:
+        """Initialise the extractor.
+
+        Raises:
+            InvalidExtractor: When the extractor is invalid.
+        """
         self.name = name = type(self).__name__
         self.logger = logging.getLogger(f"maco.extractor.{name}")
         self.logger.debug(f"initialise '{name}'")

model_setup/maco/model/model.py

@@ -29,6 +29,8 @@ class Encryption(ForbidModel):
     """Encryption usage."""
 
     class UsageEnum(str, Enum):
+        """Purpose of the encryption."""
+
         config = "config"
         communication = "communication"
         binary = "binary"
@@ -52,6 +54,8 @@ class Encryption(ForbidModel):
 
 
 class CategoryEnum(str, Enum):
+    """Category of the malware."""
+
     # Software that shows you extra promotions that you cannot control as you use your PC.
     # You wouldn't see the extra ads if you didn't have adware installed.
     adware = "adware"
@@ -274,6 +278,8 @@ class ExtractorModel(ForbidModel):
         """Binary data extracted by decoder."""
 
         class TypeEnum(str, Enum):
+            """Type of binary data."""
+
             payload = "payload"  # contained within the original file
             config = "config"  # sometimes malware uses json/formatted text for config
             other = "other"
@@ -289,6 +295,8 @@ class ExtractorModel(ForbidModel):
         # convenience for ret.encryption.append(ret.Encryption(*properties))
         # Define as class as only way to allow for this to be accessed and not have pydantic try to parse it.
         class Encryption(Encryption):
+            """Encryption usage."""
+
             pass
 
         encryption: Union[List[Encryption], Encryption, None] = None  # encryption information for the binary
@@ -383,6 +391,18 @@ class ExtractorModel(ForbidModel):
 
     proxy: List[Proxy] = []
 
+    class ICMP(ForbidModel):
+        """Usage of ICMP."""
+
+        type: Optional[int] = None
+        code: Optional[int] = None
+        header: Optional[str] = None  # Some malware uses non-standard header fields
+        hostname: Optional[str] = None
+
+        usage: Optional[ConnUsageEnum] = None
+
+    icmp: List[ICMP] = []
+
     #
     # inter process communication (IPC)
     #
@@ -436,6 +456,8 @@ class ExtractorModel(ForbidModel):
         """Direct usage of DNS."""
 
         class RecordTypeEnum(str, Enum):
+            """DNS record types."""
+
             A = "A"
             AAAA = "AAAA"
             AFSDB = "AFSDB"
@@ -512,6 +534,8 @@ class ExtractorModel(ForbidModel):
     # convenience for ret.encryption.append(ret.Encryption(*properties))
     # Define as class as only way to allow for this to be accessed and not have pydantic try to parse it.
     class Encryption(Encryption):
+        """Encryption usage."""
+
         pass
 
     encryption: List[Encryption] = []
@@ -530,6 +554,8 @@ class ExtractorModel(ForbidModel):
         """Cryptocoin usage (ransomware/miner)."""
 
         class UsageEnum(str, Enum):
+            """Cryptocoin usage."""
+
             ransomware = "ransomware"  # request money to unlock
             miner = "miner"  # use gpu/cpu to mint coins
             other = "other"
@@ -543,7 +569,11 @@ class ExtractorModel(ForbidModel):
     cryptocurrency: List[Cryptocurrency] = []
 
     class Path(ForbidModel):
+        """Path used by malware."""
+
         class UsageEnum(str, Enum):
+            """Purpose of the path."""
+
             c2 = "c2"  # file/folder issues commands to malware
             config = "config"  # config is loaded from this path
             install = "install"  # install directory/filename for malware
@@ -559,7 +589,11 @@ class ExtractorModel(ForbidModel):
     paths: List[Path] = []  # files/directories used by malware
 
     class Registry(ForbidModel):
+        """Registry usage by malware."""
+
         class UsageEnum(str, Enum):
+            """Registry usage."""
+
             persistence = "persistence"  # stay alive
             store_data = "store_data"  # generated encryption keys or config
             store_payload = "store_payload"  # malware hidden in registry key

model_setup/maco/utils.py

@@ -1,7 +1,6 @@
-# Common utilities shared between the MACO collector and configextractor-py
+"""Common utilities shared between the MACO collector and configextractor-py."""
+
 import importlib
-import importlib.machinery
-import importlib.util
 import inspect
 import json
 import logging
@@ -12,8 +11,9 @@ import shutil
 import subprocess
 import sys
 import tempfile
+from importlib.machinery import SourceFileLoader
 
-from multiprocess import Queue
+from multiprocess import Process, Queue
 
 from maco import yara
 
@@ -26,15 +26,14 @@ from base64 import b64decode
 from copy import deepcopy
 from glob import glob
 from logging import Logger
-from pkgutil import walk_packages
 from types import ModuleType
-from typing import Callable, Dict, List, Set, Tuple, Union
+from typing import Callable, Dict, List, Tuple, Union
 
 from uv import find_uv_bin
 
 from maco import model
-from maco.extractor import Extractor
 from maco.exceptions import AnalysisAbortedException
+from maco.extractor import Extractor
 
 logger = logging.getLogger("maco.lib.utils")
 
@@ -50,10 +49,14 @@ VENV_CREATE_CMD = f"{UV_BIN} venv"
 
 
 class Base64Decoder(json.JSONDecoder):
+    """JSON decoder that also base64 encodes binary data."""
+
     def __init__(self, *args, **kwargs):
+        """Initialize the decoder."""
         json.JSONDecoder.__init__(self, object_hook=self.object_hook, *args, **kwargs)
 
     def object_hook(self, obj):
+        """Hook to decode base64 encoded binary data."""  # noqa: DOC201
         if "__class__" not in obj:
             return obj
         type = obj["__class__"]
@@ -131,17 +134,38 @@ rule MACO {
 
 
 def maco_extractor_validation(module: ModuleType) -> bool:
+    """Validation function for extractors.
+
+    Returns:
+        (bool): True if extractor belongs to MACO, False otherwise.
+    """
     if inspect.isclass(module):
         # 'author' has to be implemented otherwise will raise an exception according to MACO
         return hasattr(module, "author") and module.author
     return False
 
 
-def maco_extract_rules(module: Extractor) -> bool:
+def maco_extract_rules(module: Extractor) -> str:
+    """Extracts YARA rules from extractor.
+
+    Returns:
+     (str): YARA rules
+    """
     return module.yara_rule
 
 
 def scan_for_extractors(root_directory: str, scanner: yara.Rules, logger: Logger) -> Tuple[List[str], List[str]]:
+    """Looks for extractors using YARA rules.
+
+    Args:
+        root_directory (str): Root directory containing extractors
+        scanner (yara.Rules): Scanner to look for extractors using YARA rules
+        logger (Logger): Logger to use
+
+    Returns:
+        Tuple[List[str], List[str]]: Returns a list of extractor directories and extractor files
+
+    """
     extractor_dirs = set([root_directory])
     extractor_files = []
 
@@ -177,17 +201,22 @@ def scan_for_extractors(root_directory: str, scanner: yara.Rules, logger: Logger
                     with open(path, "rb") as f:
                         data = f.read()
 
-                    with open(path, "wb") as f:
-                        # Replace any relative importing with absolute
-                        curr_dir = os.path.dirname(path)
-                        split = curr_dir.split("/")[::-1]
-                        for pattern in [RELATIVE_FROM_IMPORT_RE, RELATIVE_FROM_RE]:
-                            for match in pattern.findall(data):
-                                depth = match.count(b".")
-                                abspath = ".".join(split[depth - 1 : split.index(package) + 1][::-1])
-                                abspath += "." if pattern == RELATIVE_FROM_RE else ""
-                                data = data.replace(f"from {match.decode()}".encode(), f"from {abspath}".encode(), 1)
-                        f.write(data)
+                    # Replace any relative importing with absolute
+                    changed_imports = False
+                    curr_dir = os.path.dirname(path)
+                    split = curr_dir.split("/")[::-1]
+                    for pattern in [RELATIVE_FROM_IMPORT_RE, RELATIVE_FROM_RE]:
+                        for match in pattern.findall(data):
+                            depth = match.count(b".")
+                            abspath = ".".join(split[depth - 1 : split.index(package) + 1][::-1])
+                            abspath += "." if pattern == RELATIVE_FROM_RE else ""
+                            data = data.replace(f"from {match.decode()}".encode(), f"from {abspath}".encode(), 1)
+                            changed_imports = True
+
+                    # only write extractor files if imports were changed
+                    if changed_imports:
+                        with open(path, "wb") as f:
+                            f.write(data)
 
                 if scanner.match(path):
                     # Add directory to list of hits for venv creation
@@ -282,7 +311,16 @@ def _install_required_packages(create_venv: bool, directories: List[str], python
     return venvs
 
 
-def find_and_insert_venv(path: str, venvs: List[str]):
+def find_and_insert_venv(path: str, venvs: List[str]) -> Tuple[str, str]:
+    """Finds the closest virtual environment to the extractor and inserts it into the PATH.
+
+    Args:
+        path (str): Path of extractor
+        venvs (List[str]): List of virtual environments
+
+    Returns:
+        (Tuple[str, str]): Virtual environment and site-packages path that's closest to the extractor
+    """
     venv = None
     for venv in sorted(venvs, reverse=True):
         venv_parent = os.path.dirname(venv)
@@ -303,14 +341,53 @@ def find_and_insert_venv(path: str, venvs: List[str]):
     return venv, site_package
 
 
+def register_extractor_module(
+    extractor_source_file: str,
+    module_name: str,
+    venvs: List[str],
+    extractor_module_callback: Callable[[ModuleType, str], None],
+    logger: Logger,
+):
+    """Register the extractor module in isolation.
+
+    Args:
+        extractor_source_file (str): Path to source file of extractor
+        module_name (str): The name of the module relative to the package directory
+        venvs (List[str]): List of virtual environments
+        extractor_module_callback (Callable[[ModuleType, str], None]): Callback used to register extractors
+        logger (Logger): Logger to use
+
+    """
+    try:
+        logger.info(f"Inspecting '{extractor_source_file}' for extractors..")
+        venv, site_packages = find_and_insert_venv(extractor_source_file, venvs)
+        loader = SourceFileLoader(
+            module_name,
+            extractor_source_file,
+        )
+        extractor_module_callback(loader.load_module(), venv)
+    finally:
+        # Cleanup virtual environment that was loaded into PATH
+        if venv and site_packages in sys.path:
+            sys.path.remove(site_packages)
+
+
 def register_extractors(
     current_directory: str,
     venvs: List[str],
     extractor_files: List[str],
     extractor_module_callback: Callable[[ModuleType, str], None],
     logger: Logger,
-    default_loaded_modules: Set[str] = set(sys.modules.keys()),
 ):
+    """Register extractors with in the current directory.
+
+    Args:
+        current_directory (str): Current directory to register extractors found
+        venvs (List[str]): List of virtual environments
+        extractor_files (List[str]): List of extractor files found
+        extractor_module_callback (Callable[[ModuleType, str], None]): Callback used to register extractors
+        logger (Logger): Logger to use
+    """
     package_name = os.path.basename(current_directory)
     parent_directory = os.path.dirname(current_directory)
     if venvs and package_name in sys.modules:
@@ -325,74 +402,25 @@ def register_extractors(
         sys.path.insert(1, current_directory)
         sys.path.insert(1, parent_directory)
 
-        # Insert any virtual environment necessary to load directory as package
-        package_venv, package_site_packages = find_and_insert_venv(current_directory, venvs)
-        package = importlib.import_module(package_name)
+        # Load the potential extractors directly from the source file
+        registration_processes = []
+        for extractor_source_file in extractor_files:
+            module_name = extractor_source_file.replace(f"{parent_directory}/", "").replace("/", ".")[:-3]
+            p = Process(
+                target=register_extractor_module,
+                args=(extractor_source_file, module_name, venvs, extractor_module_callback, logger),
+            )
+            p.start()
+            registration_processes.append(p)
 
-        # Walk through our new package and find the extractors that YARA identified
-        for module_path, module_name, ispkg in walk_packages(package.__path__, package.__name__ + "."):
-            if ispkg:
-                # Skip packages
-                continue
+        for p in registration_processes:
+            p.join()
 
-            module_path = os.path.realpath(os.path.join(module_path.path, module_name.rsplit(".", 1)[1]) + ".py")
-            if module_path in extractor_files:
-                # Cross this extractor off the list of extractors to find
-                logger.debug(f"Inspecting '{module_name}' for extractors..")
-                extractor_files.remove(module_path)
-                try:
-                    # This is an extractor we've been looking for, load the module and invoke callback
-                    venv, site_packages = find_and_insert_venv(module_path, venvs)
-                    module = importlib.import_module(module_name)
-                    module.__file__ = os.path.realpath(module.__file__)
-
-                    # Patch the original directory information into the module
-                    original_package_name = os.path.basename(current_directory)
-                    module.__name__ = module.__name__.replace(package_name, original_package_name)
-                    module.__package__ = module.__package__.replace(package_name, original_package_name)
-                    extractor_module_callback(module, venv)
-                finally:
-                    # Cleanup virtual environment that was loaded into PATH
-                    if venv and site_packages in sys.path:
-                        sys.path.remove(site_packages)
-
-            if not extractor_files:
-                return
     finally:
         # Cleanup changes made to PATH
         sys.path.remove(parent_directory)
         sys.path.remove(current_directory)
 
-        if package_venv and package_site_packages in sys.path:
-            sys.path.remove(package_site_packages)
-
-        # Remove any modules that were loaded to deconflict with later modules loads
-        [sys.modules.pop(k) for k in set(sys.modules.keys()) - default_loaded_modules]
-
-    # If there still exists extractor files we haven't found yet, try searching in the available subdirectories
-    if extractor_files:
-        for dir in os.listdir(current_directory):
-            path = os.path.join(current_directory, dir)
-            if dir == "__pycache__":
-                # Ignore the cache created
-                continue
-            elif dir.endswith(".egg-info"):
-                # Ignore these directories
-                continue
-            elif dir.startswith("."):
-                # Ignore hidden directories
-                continue
-
-            if os.path.isdir(path):
-                # Check subdirectory to find the rest of the detected extractors
-                register_extractors(
-                    path, venvs, extractor_files, extractor_module_callback, logger, default_loaded_modules
-                )
-
-            if not extractor_files:
-                # We were able to find all the extractor files
-                break
-
 
 def proxy_logging(queue: Queue, callback: Callable[[ModuleType, str], None], *args, **kwargs):
     """Ensures logging is set up correctly for a child process and then executes the callback."""
@@ -413,6 +441,17 @@ def import_extractors(
     python_version: str = f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}",
     skip_install: bool = False,
 ):
+    """Import extractors in a given directory.
+
+    Args:
+        extractor_module_callback (Callable[[ModuleType, str], bool]): Callback used to register extractors
+        root_directory (str): Root directory to look for extractors
+        scanner (yara.Rules): Scanner to look for extractors that match YARA rule
+        create_venv (bool): Create/Use virtual environments
+        logger (Logger): Logger to use
+        python_version (str): Version of python to use when creating virtual environments
+        skip_install (bool): Skip installation of Python dependencies for extractors
+    """
     extractor_dirs, extractor_files = scan_for_extractors(root_directory, scanner, logger)
 
     logger.info(f"Extractor files found based on scanner ({len(extractor_files)}).")
@@ -448,7 +487,24 @@ def run_extractor(
     venv_script=VENV_SCRIPT,
     json_decoder=Base64Decoder,
 ) -> Union[Dict[str, dict], model.ExtractorModel]:
-    """Runs the maco extractor against sample either in current process or child process."""
+    """Runs the maco extractor against sample either in current process or child process.
+
+    Args:
+        sample_path (str): Path to sample
+        module_name (str): Name of extractor module
+        extractor_class (str): Name of extractor class in module
+        module_path (str): Path to Python module containing extractor
+        venv (str): Path to virtual environment associated to extractor
+        venv_script (str): Script to run extractor in a virtual environment
+        json_decoder (Base64Decoder): Decoder used for JSON
+
+    Raises:
+        AnalysisAbortedException: Raised when extractor voluntarily terminates execution
+        Exception: Raised when extractor raises an exception
+
+    Returns:
+        Union[Dict[str, dict], model.ExtractorModel]: Results from extractor
+    """
     if not venv:
         key = f"{module_name}_{extractor_class}"
         if key not in _loaded_extractors:

model_setup/maco/yara.py

@@ -1,25 +1,34 @@
+"""yara-python facade that uses yara-x."""
+
 import re
 from collections import namedtuple
 from itertools import cycle
-from typing import Dict
+from typing import Dict, List
 
 import yara_x
 
-RULE_ID_RE = re.compile("(\w+)? ?rule (\w+)")
-
+from maco.exceptions import SyntaxError
 
-class SyntaxError(Exception): ...
+RULE_ID_RE = re.compile("(\w+)? ?rule (\w+)")
 
 
 # Create interfaces that resembles yara-python (but is running yara-x under the hood)
 class StringMatchInstance:
+    """Instance of a string match."""
+
     def __init__(self, match: yara_x.Match, file_content: bytes):
+        """Initializes StringMatchInstance."""
         self.matched_data = file_content[match.offset : match.offset + match.length]
         self.matched_length = match.length
         self.offset = match.offset
         self.xor_key = match.xor_key
 
     def plaintext(self) -> bytes:
+        """Plaintext of the matched data.
+
+        Returns:
+            (bytes): Plaintext of the matched cipher text
+        """
         if not self.xor_key:
             # No need to XOR the matched data
             return self.matched_data
@@ -28,17 +37,28 @@ class StringMatchInstance:
 
 
 class StringMatch:
+    """String match."""
+
     def __init__(self, pattern: yara_x.Pattern, file_content: bytes):
+        """Initializes StringMatch."""
         self.identifier = pattern.identifier
         self.instances = [StringMatchInstance(match, file_content) for match in pattern.matches]
         self._is_xor = any([match.xor_key for match in pattern.matches])
 
     def is_xor(self):
+        """Checks if string match is xor'd.
+
+        Returns:
+            (bool): True if match is xor'd
+        """
         return self._is_xor
 
 
 class Match:
+    """Match."""
+
     def __init__(self, rule: yara_x.Rule, file_content: bytes):
+        """Initializes Match."""
         self.rule = rule.identifier
         self.namespace = rule.namespace
         self.tags = list(rule.tags) or []
@@ -50,7 +70,14 @@ class Match:
 
 
 class Rules:
+    """Rules."""
+
     def __init__(self, source: str = None, sources: Dict[str, str] = None):
+        """Initializes Rules.
+
+        Raises:
+            SyntaxError: Raised when there's a syntax error in the YARA rule.
+        """
         Rule = namedtuple("Rule", "identifier namespace is_global")
         if source:
             sources = {"default": source}
@@ -69,10 +96,20 @@ class Rules:
             raise SyntaxError(e)
 
     def __iter__(self):
+        """Iterate over rules.
+
+        Yields:
+            YARA rules
+        """
         for rule in self._rules:
             yield rule
 
-    def match(self, filepath: str = None, data: bytes = None):
+    def match(self, filepath: str = None, data: bytes = None) -> List[Match]:
+        """Performs a scan to check for YARA rules matches based on the file, either given by path or buffer.
+
+        Returns:
+            (List[Match]): A list of YARA matches.
+        """
         if filepath:
             with open(filepath, "rb") as fp:
                 data = fp.read()
@@ -81,4 +118,9 @@ class Rules:
 
 
 def compile(source: str = None, sources: Dict[str, str] = None) -> Rules:
+    """Compiles YARA rules from source or from sources.
+
+    Returns:
+        (Rules): a Rules object
+    """
     return Rules(source, sources)