Merge pull request #71 from seb-acsc/feat/performance

Improve non-venv performance by 85x

demo_extractors/complex/complex.py

@@ -1,9 +1,8 @@
 from io import BytesIO
 from typing import List, Optional
 
-from maco import extractor, model, yara
-
 from demo_extractors.complex import complex_utils
+from maco import extractor, model, yara
 
 
 class Complex(extractor.Extractor):

demo_extractors/limit_other.py

@@ -1,9 +1,8 @@
 from io import BytesIO
 from typing import Dict, List, Optional
 
-from maco import extractor, model, yara
-
 from demo_extractors import shared
+from maco import extractor, model, yara
 
 
 class LimitOther(extractor.Extractor):
@@ -24,6 +23,10 @@ class LimitOther(extractor.Extractor):
         """
 
     def run(self, stream: BytesIO, matches: List[yara.Match]) -> Optional[model.ExtractorModel]:
+        # import httpx at runtime so we can test that requirements.txt is installed dynamically without breaking
+        # the tests that do direct importing
+        import httpx
+
         # use a custom model that inherits from ExtractorModel
         # this model defines what can go in the 'other' dict
         tmp = shared.MyCustomModel(family="specify_other")

demo_extractors/requirements.txt

@@ -0,0 +1 @@
+httpx

demo_extractors/shared.py

@@ -1,4 +1,5 @@
 from typing import Optional
+
 import pydantic
 
 from maco import model

maco/base_test.py

@@ -32,14 +32,19 @@ class BaseTest(unittest.TestCase):
     # I recommend something like os.path.join(__file__, "../../extractors")
     # if your extractors are in a folder 'extractors' next to a folder of tests
     path: str = None
+    create_venv: bool=False
 
-    def setUp(self) -> None:
-        if not self.name or not self.path:
+    @classmethod
+    def setUpClass(cls) -> None:
+        if not cls.name or not cls.path:
             raise Exception("name and path must be set")
-        self.c = collector.Collector(self.path, include=[self.name])
+        cls.c = collector.Collector(cls.path, include=[cls.name], create_venv=cls.create_venv)
+        return super().setUpClass()
+
+    def test_default_metadata(self):
+        """Require extractor to be loadable and valid."""
         self.assertIn(self.name, self.c.extractors)
         self.assertEqual(len(self.c.extractors), 1)
-        return super().setUp()
 
     def extract(self, stream):
         """Return results for running extractor over stream, including yara check."""
@@ -49,18 +54,20 @@ def extract(self, stream):
         resp = self.c.extract(stream, self.name)
         return resp
 
-    def _get_location(self) -> str:
+    @classmethod
+    def _get_location(cls) -> str:
         """Return path to child class that implements this class."""
         # import child module
-        module = type(self).__module__
+        module = cls.__module__
         i = importlib.import_module(module)
         # get location to child module
         return i.__file__
 
-    def load_cart(self, filepath: str) -> io.BytesIO:
+    @classmethod
+    def load_cart(cls, filepath: str) -> io.BytesIO:
         """Load and unneuter a test file (likely malware) into memory for processing."""
         # it is nice if we can load files relative to whatever is implementing base_test
-        dirpath = os.path.split(self._get_location())[0]
+        dirpath = os.path.split(cls._get_location())[0]
         # either filepath is absolute, or should be loaded relative to child of base_test
         filepath = os.path.join(dirpath, filepath)
         if not os.path.isfile(filepath):

maco/cli.py

@@ -179,7 +179,7 @@ def main():
     parser.add_argument(
         "--create_venv",
         action="store_true",
-        help="Creates venvs for every requirements.txt found (only applies when extractor path is a directory)",
+        help="Creates venvs for every requirements.txt found (only applies when extractor path is a directory). This runs much slower than the alternative but may be necessary when there are many extractors with conflicting dependencies.",
     )
     args = parser.parse_args()
     inc = args.include.split(",") if args.include else []

maco/extractor.py

@@ -51,14 +51,14 @@ def __init__(self) -> None:
         # check yara rules conform to expected structure
         # we throw away these compiled rules as we need all rules in system compiled together
         try:
-            rules = yara.compile(source=self.yara_rule)
+            self.yara_compiled = yara.compile(source=self.yara_rule)
         except yara.SyntaxError as e:
             raise InvalidExtractor(f"{self.name} - invalid yara rule") from e
         # need to track which plugin owns the rules
-        self.yara_rule_names = [x.identifier for x in rules]
-        if not len(list(rules)):
+        self.yara_rule_names = [x.identifier for x in self.yara_compiled]
+        if not len(list(self.yara_compiled)):
             raise InvalidExtractor(f"{name} must define at least one yara rule")
-        for x in rules:
+        for x in self.yara_compiled:
             if x.is_global:
                 raise InvalidExtractor(f"{x.identifier} yara rule must not be global")