Allow to extend class and reuse original main (#23)

karton/yaramatcher/__version__.py

@@ -1 +1 @@
-__version__ = "1.2.0"
+__version__ = "1.3.0"

karton/yaramatcher/yaramatcher.py

@@ -1,3 +1,4 @@
+import argparse
 import logging
 import os
 import re
@@ -6,7 +7,7 @@
 from typing import List, Optional
 
 import yara  # type: ignore
-from karton.core import Config, Karton, Task  # type: ignore
+from karton.core import Config, Karton, Task
 
 from .__version__ import __version__
 
@@ -31,7 +32,6 @@ class YaraHandler:
     """
 
     def __init__(self, path: Optional[str] = None) -> None:
-        super().__init__()
         # load and compile all Yara rules from a folder
         yara_path = path or "rules"
         rule_paths = []
@@ -97,17 +97,15 @@ def args_parser(cls):
         return parser
 
     @classmethod
-    def main(cls):
-        parser = cls.args_parser()
-        args = parser.parse_args()
+    def config_from_args(cls, config: Config, args: argparse.Namespace) -> None:
+        super().config_from_args(config, args)
+        config.load_from_dict({"yaramatcher": {"rules": args.rules}})
 
-        config = Config(args.config_file)
-        service = YaraMatcher(config=config, yara_rule_dir=args.rules)
-        service.loop()
-
-    def __init__(self, yara_rule_dir: Optional[str] = None, *args, **kwargs) -> None:
+    def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
-        self.yara_handler = YaraHandler(path=yara_rule_dir or "rules")
+        self.yara_handler = YaraHandler(
+            path=self.config.get("yaramatcher", "rules", fallback="rules")
+        )
 
     def scan_sample(self, sample: bytes) -> List[str]:
         # Get all matches for this sample

tests/test_matching.py

@@ -1,9 +1,15 @@
 import unittest
 from karton.yaramatcher import YaraMatcher, normalize_rule_name
-from karton.core.test import KartonTestCase
+from karton.core.test import KartonTestCase, ConfigMock
 from karton.core import Resource, Task
 
 
+class YaraMatcherConfigMock(ConfigMock):
+    def __init__(self):
+        super().__init__()
+        self._config["yaramatcher"] = {"rules": "tests/testdata/rules"}
+
+
 class TestUtils(unittest.TestCase):
     def test_normalize_rule_name(self) -> None:
         self.assertEqual("win_remcos", normalize_rule_name("win_remcos_auto"))
@@ -14,7 +20,7 @@ def test_normalize_rule_name(self) -> None:
 
 class YaraMatcherTestBasic(KartonTestCase):
     karton_class = YaraMatcher
-    kwargs = {"yara_rule_dir": "tests/testdata/rules"}
+    config = YaraMatcherConfigMock()
 
     def test_pass(self) -> None:
         res = Resource("sample", b"z")