diff --git a/karton/yaramatcher/__version__.py b/karton/yaramatcher/__version__.py index c68196d..67bc602 100644 --- a/karton/yaramatcher/__version__.py +++ b/karton/yaramatcher/__version__.py @@ -1 +1 @@ -__version__ = "1.2.0" +__version__ = "1.3.0" diff --git a/karton/yaramatcher/yaramatcher.py b/karton/yaramatcher/yaramatcher.py index 8f32f92..15638d3 100644 --- a/karton/yaramatcher/yaramatcher.py +++ b/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 diff --git a/tests/test_matching.py b/tests/test_matching.py index 26ff92e..a827c7a 100644 --- a/tests/test_matching.py +++ b/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")