Merge pull request #79 from CybercentreCanada/general_improvements

General improvements

demo_extractors/requirements.txt

@@ -1 +1,4 @@
 httpx
+
+# Install maco from source for testing
+../

demo_extractors/terminator.py

@@ -0,0 +1,17 @@
+from io import BytesIO
+from typing import List, Optional
+
+from maco import extractor, model, yara
+from maco.exceptions import AnalysisAbortedException
+
+
+class Terminator(extractor.Extractor):
+    """Terminates early during extraction"""
+
+    family = "terminator"
+    author = "skynet"
+    last_modified = "1997-08-29"
+
+    def run(self, stream: BytesIO, matches: List[yara.Match]) -> Optional[model.ExtractorModel]:
+        # Terminate early and indicate I can't run on this sample
+        raise AnalysisAbortedException("I can't run on this sample")

maco/cli.py

@@ -3,18 +3,20 @@
 import argparse
 import base64
 import binascii
+import cart
 import hashlib
 import io
 import json
 import logging
 import os
 import sys
-from typing import BinaryIO, List, Tuple
 
-import cart
+from importlib.metadata import version
+from typing import BinaryIO, List, Tuple
 
 from maco import collector
 
+
 logger = logging.getLogger("maco.lib.cli")
 
 
@@ -199,6 +201,13 @@ def main():
         action="store_true",
         help="Force installation of Python dependencies for extractors (in both host and virtual environments).",
     )
+    parser.add_argument(
+        "--version",
+        action="version",
+        version=f"version: {version('maco')}",
+        help="Show version of MACO",
+    )
+
     args = parser.parse_args()
     inc = args.include.split(",") if args.include else []
     exc = args.exclude.split(",") if args.exclude else []

maco/collector.py

@@ -13,6 +13,7 @@
 from pydantic import BaseModel
 
 from maco import extractor, model, utils, yara
+from maco.exceptions import AnalysisAbortedException
 
 
 class ExtractorLoadError(Exception):
@@ -191,6 +192,9 @@ def extract(
                         venv=extractor["venv"],
                     )
                 )
+        except AnalysisAbortedException:
+            # Extractor voluntarily aborted analysis of sample
+            return
         except Exception:
             # caller can deal with the exception
             raise

maco/exceptions.py

@@ -0,0 +1,3 @@
+# 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): ...

maco/utils.py

@@ -34,6 +34,7 @@
 
 from maco import model
 from maco.extractor import Extractor
+from maco.exceptions import AnalysisAbortedException
 
 logger = logging.getLogger("maco.lib.utils")
 
@@ -514,9 +515,15 @@ def run_extractor(
                     exception = stderr
                     if delim in exception:
                         exception = f"{delim}{exception.split(delim, 1)[1]}"
-                    # print extractor logging at error level
-                    logger.error(f"maco extractor raised exception, stderr:\n{stderr}")
-                    raise Exception(exception) from e
+                    if "maco.exceptions.AnalysisAbortedException" in exception:
+                        # Extractor voluntarily terminated, re-raise exception to be handled by collector
+                        raise AnalysisAbortedException(
+                            exception.split("maco.exceptions.AnalysisAbortedException: ")[-1]
+                        )
+                    else:
+                        # print extractor logging at error level
+                        logger.error(f"maco extractor raised exception, stderr:\n{stderr}")
+                        raise Exception(exception) from e
                 # ensure that extractor logging is available
                 logger.info(f"maco extractor stderr:\n{stderr}")
     return loaded

tests/test_base_test.py

@@ -61,3 +61,29 @@ def test_extract(self):
         ret = self.extract(self.load_cart("data/trigger_complex.txt.cart"))
         self.assertEqual(ret["family"], "complex")
         self.assertEqual(ret["version"], "5")
+
+
+class TestTerminator(base_test.BaseTest):
+    """Test that terminator extractor can be used in base environment."""
+
+    name = "Terminator"
+    path = os.path.join(__file__, "../../demo_extractors")
+    create_venv = False
+
+    def test_extract(self):
+        """Tests that we can run an extractor through maco."""
+        ret = self.extract(self.load_cart("data/trigger_complex.txt.cart"))
+        self.assertEqual(ret, None)
+
+
+class TestTerminatorVenv(base_test.BaseTest):
+    """Test that terminator extractor can be used in base environment."""
+
+    name = "Terminator"
+    path = os.path.join(__file__, "../../demo_extractors")
+    create_venv = True
+
+    def test_extract(self):
+        """Tests that we can run an extractor through maco."""
+        ret = self.extract(self.load_cart("data/trigger_complex.txt.cart"))
+        self.assertEqual(ret, None)

tests/test_demo_extractors.py

@@ -11,12 +11,7 @@ def test_complex(self):
         collector = Collector(os.path.join(__file__, "../../demo_extractors"))
         self.assertEqual(
             set(collector.extractors.keys()),
-            {
-                "Elfy",
-                "Nothing",
-                "Complex",
-                "LimitOther",
-            },
+            {"Elfy", "Nothing", "Complex", "LimitOther", "Terminator"},
         )
 
         with open(path_file, "rb") as stream:

tests/test_helpers.py

@@ -25,14 +25,7 @@ def test_compile_yara(self):
         m = collector.Collector(target)
         self.assertEqual(
             {x.identifier for x in m.rules},
-            {
-                "Elfy",
-                "Complex",
-                "ComplexSubtext",
-                "Nothing",
-                "ComplexAlt",
-                "LimitOther",
-            },
+            {"Elfy", "Complex", "ComplexSubtext", "Nothing", "ComplexAlt", "LimitOther", "Terminator"},
         )