From ecff817b492efab4e255cd609fd2cdc5041fa2c9 Mon Sep 17 00:00:00 2001 From: wunder957 Date: Fri, 22 Sep 2023 16:44:03 +0800 Subject: [PATCH] Feature: template for tracer (#76) * Support include_extension and init template * Suuport ShTracer template * Add example for tracer template --- README.md | 2 +- docs/source/managers/tracer.rst | 7 ++ duetector/cli/main.py | 1 - duetector/collectors/models.py | 1 + duetector/managers/analyzer.py | 5 +- duetector/managers/base.py | 8 +- duetector/managers/collector.py | 5 +- duetector/managers/filter.py | 5 +- duetector/managers/tracer.py | 107 +++++++++++++++++- duetector/monitors/sh_monitor.py | 3 +- duetector/static/config.toml | 5 + duetector/tools/config_generator.py | 2 +- duetector/tracers/base.py | 2 +- examples/README.md | 1 + .../config/sh-tracer-template.config.toml | 9 ++ tests/test_bcc_monitor.py | 4 +- tests/test_tracer_template.py | 45 ++++++++ 17 files changed, 193 insertions(+), 19 deletions(-) create mode 100644 examples/config/sh-tracer-template.config.toml create mode 100644 tests/test_tracer_template.py diff --git a/README.md b/README.md index f58eb19..703eca3 100644 --- a/README.md +++ b/README.md @@ -40,7 +40,7 @@ Try simple user case: [Simplest Open Count](./docs/usercases/simplest-open-count ## Feature -- Plug-in system support, see [examples](./examples/) for more details +- Plug-in system support, see [examples](./examples/README.md) for more details - [X] Custom `Tracer` and `TracerManager` - [X] Custom `Filters` and `FilterManager` - [X] Custom `Collector` and `CollectorManager` diff --git a/docs/source/managers/tracer.rst b/docs/source/managers/tracer.rst index 3c53fc6..8c2c29e 100644 --- a/docs/source/managers/tracer.rst +++ b/docs/source/managers/tracer.rst @@ -11,3 +11,10 @@ TracerManager :undoc-members: :inherited-members: :show-inheritance: + + +.. autoclass:: duetector.managers.tracer.TracerTemplate + :members: + :undoc-members: + :inherited-members: + :show-inheritance: diff --git a/duetector/cli/main.py b/duetector/cli/main.py index 1979d45..b54f02c 100644 --- a/duetector/cli/main.py +++ b/duetector/cli/main.py @@ -7,7 +7,6 @@ import click -from duetector.analyzer.db import DBAnalyzer from duetector.config import CONFIG_PATH, ConfigLoader from duetector.log import logger from duetector.managers.analyzer import AnalyzerManager diff --git a/duetector/collectors/models.py b/duetector/collectors/models.py index b4084fe..dfb074e 100644 --- a/duetector/collectors/models.py +++ b/duetector/collectors/models.py @@ -79,6 +79,7 @@ def from_namedtuple(tracer, data: NamedTuple) -> Tracking: # type: ignore # Is instance of tracer tracer_name = getattr(tracer, "name", tracer.__class__.__name__) + tracer_name = tracer_name.lower() args = { "tracer": tracer_name, "extended": {}, diff --git a/duetector/managers/analyzer.py b/duetector/managers/analyzer.py index 2d58208..351ae6d 100644 --- a/duetector/managers/analyzer.py +++ b/duetector/managers/analyzer.py @@ -39,10 +39,11 @@ def __init__(self, config: Optional[Dict[str, Any]] = None, *args, **kwargs): self.pm = pluggy.PluginManager(PROJECT_NAME) self.pm.add_hookspecs(sys.modules[__name__]) - self.pm.load_setuptools_entrypoints(PROJECT_NAME) + if self.include_extension: + self.pm.load_setuptools_entrypoints(PROJECT_NAME) self.register(duetector.analyzer.register) - def init(self, analyzer_type=Analyzer, ignore_disabled=True) -> List[Analyzer]: + def init(self, analyzer_type=Analyzer, ignore_disabled=True, *args, **kwargs) -> List[Analyzer]: """ Initialize all analyzers from config. diff --git a/duetector/managers/base.py b/duetector/managers/base.py index 1255bd9..6446fd2 100644 --- a/duetector/managers/base.py +++ b/duetector/managers/base.py @@ -1,4 +1,4 @@ -from typing import Any, Dict, Optional +from typing import Any, Dict, List, Optional import pluggy @@ -64,3 +64,9 @@ def disabled(self): If current manager is disabled. """ return self.config.disabled + + def init(self, *args, **kwargs) -> List[Any]: + """ + Initialize manager. + """ + raise NotImplementedError diff --git a/duetector/managers/collector.py b/duetector/managers/collector.py index d35f1fc..9b93ca4 100644 --- a/duetector/managers/collector.py +++ b/duetector/managers/collector.py @@ -41,10 +41,11 @@ def __init__(self, config: Optional[Dict[str, Any]] = None, *args, **kwargs): self.pm = pluggy.PluginManager(PROJECT_NAME) self.pm.add_hookspecs(sys.modules[__name__]) - self.pm.load_setuptools_entrypoints(PROJECT_NAME) + if self.include_extension: + self.pm.load_setuptools_entrypoints(PROJECT_NAME) self.register(duetector.collectors.register) - def init(self, ignore_disabled=True) -> List[Collector]: + def init(self, ignore_disabled=True, *args, **kwargs) -> List[Collector]: """ Initialize all collectors from config. diff --git a/duetector/managers/filter.py b/duetector/managers/filter.py index 0865c4a..a38031b 100644 --- a/duetector/managers/filter.py +++ b/duetector/managers/filter.py @@ -39,10 +39,11 @@ def __init__(self, config: Optional[Dict[str, Any]] = None, *args, **kwargs): self.pm = pluggy.PluginManager(PROJECT_NAME) self.pm.add_hookspecs(sys.modules[__name__]) - self.pm.load_setuptools_entrypoints(PROJECT_NAME) + if self.include_extension: + self.pm.load_setuptools_entrypoints(PROJECT_NAME) self.register(duetector.filters.register) - def init(self, ignore_disabled=True) -> List[Filter]: + def init(self, ignore_disabled=True, *args, **kwargs) -> List[Filter]: """ Initialize all filters from config. diff --git a/duetector/managers/tracer.py b/duetector/managers/tracer.py index dfb1ca5..e088a09 100644 --- a/duetector/managers/tracer.py +++ b/duetector/managers/tracer.py @@ -1,13 +1,14 @@ import sys -from typing import Any, Dict, List, Optional +from typing import Any, Dict, Iterable, List, Optional, Union import pluggy import duetector.tracers.register +from duetector.config import Config, Configuable from duetector.extension.tracer import project_name from duetector.log import logger from duetector.managers.base import Manager -from duetector.tracers.base import Tracer +from duetector.tracers.base import ShellTracer, Tracer PROJECT_NAME = project_name #: Default project name for pluggy hookspec = pluggy.HookspecMarker(PROJECT_NAME) @@ -22,6 +23,82 @@ def init_tracer(config) -> Optional[Tracer]: """ +class TracerTemplate(Configuable): + """ + Using template to generate tracers. + + Tracers are initialized from config, and can be ``disabled`` by config. + Tracer type is defined by ``_avaliable_tracer_type``. + + + Example: + + .. code-block:: toml + + [tracer.template.sh] + pstracer = { "comm" = ["ps", "-aux"], config = { "enable_cache" = false } } + """ + + _avaliable_tracer_type = { + "sh": ShellTracer, + } + """ + Available tracer type. + """ + + default_config = { + "disabled": False, + **{k: dict() for k in _avaliable_tracer_type.keys()}, + } + """ + Default config for ``TracerTemplate``. + """ + + @property + def disabled(self) -> bool: + """ + Whether the template is disabled. + """ + return self.config.disabled + + def __init__(self, config: Optional[Dict[str, Any]] = None, *args, **kwargs): + super().__init__(config, *args, **kwargs) + + def _get_wrap_tracer( + self, tracer_type: Union[str, type], tracer_name: str, kwargs: Config + ) -> Tracer: + """ + Get a wrapper class for tracer type. + """ + if isinstance(tracer_type, str): + tracer_type = self._avaliable_tracer_type[tracer_type] + + class WrapTracer(tracer_type): + config_scope = None + + tracer_config = kwargs.pop("config", {}) + WrapTracer.__name__ = WrapTracer.name = tracer_name + for k, v in kwargs.items(): + setattr(WrapTracer, k, v) + + return WrapTracer(tracer_config) + + def init(self) -> List[Tracer]: + """ + Initialize all tracers from config. + """ + if self.disabled: + logger.info("TracerTemplate disabled.") + return [] + + objs = [] + for k, tracer_type in self._avaliable_tracer_type.items(): + for tracer_name, kwargs in self.config._config_dict[k].items(): + objs.append(self._get_wrap_tracer(tracer_type, tracer_name, kwargs)) + + return objs + + class TracerManager(Manager): """ Manager for all tracers. @@ -29,6 +106,11 @@ class TracerManager(Manager): Tracers are initialized from config, and can be ``disabled`` by config. """ + default_config = { + **Manager.default_config, + "template": {**TracerTemplate.default_config}, + } + config_scope = "tracer" """ Config scope for ``TracerManager``. @@ -39,23 +121,38 @@ def __init__(self, config: Optional[Dict[str, Any]] = None, *args, **kwargs): self.pm = pluggy.PluginManager(PROJECT_NAME) self.pm.add_hookspecs(sys.modules[__name__]) - self.pm.load_setuptools_entrypoints(PROJECT_NAME) + if self.include_extension: + self.pm.load_setuptools_entrypoints(PROJECT_NAME) self.register(duetector.tracers.register) - def init(self, tracer_type=Tracer, ignore_disabled=True) -> List[Tracer]: + self.template = TracerTemplate(self.config.template) + + def init( + self, + tracer_type=Tracer, + ignore_disabled=True, + include_template=True, + *args, + **kwargs, + ) -> List[Tracer]: """ Initialize all tracers from config. Args: tracer_type: Only return tracers of this type ignore_disabled: Ignore disabled tracers + include_template: Include tracers from template, ``False`` when used to generate configuration. """ if self.disabled: logger.info("TracerManager disabled.") return [] objs = [] - for f in self.pm.hook.init_tracer(config=self.config._config_dict): + tracers = list(self.pm.hook.init_tracer(config=self.config)) + if include_template: + tracers.extend(self.template.init()) + + for f in tracers: if not f or (f.disabled and ignore_disabled): logger.debug(f"Tracer {f.__class__.__name__} is not available (None or Disabled)") continue diff --git a/duetector/monitors/sh_monitor.py b/duetector/monitors/sh_monitor.py index f2fce3c..b19b818 100644 --- a/duetector/monitors/sh_monitor.py +++ b/duetector/monitors/sh_monitor.py @@ -1,4 +1,5 @@ import subprocess +from datetime import datetime from typing import Any, Callable, Dict, List, NamedTuple, Optional from duetector.collectors.base import Collector @@ -50,7 +51,7 @@ def _(): tracer.set_cache(output) callback = self.callbacks[tracer] - callback(tracer.data_t(output=output)) + callback(tracer.data_t(output=output, dt=datetime.now())) return _ diff --git a/duetector/static/config.toml b/duetector/static/config.toml index de8bb51..d1a9dae 100644 --- a/duetector/static/config.toml +++ b/duetector/static/config.toml @@ -34,6 +34,11 @@ exclude_gid = [ disabled = false include_extension = true +[tracer.template] +disabled = false + +[tracer.template.sh] + [tracer.clonetracer] disabled = false attach_event = "__x64_sys_clone" diff --git a/duetector/tools/config_generator.py b/duetector/tools/config_generator.py index 7313bd1..b2ce894 100644 --- a/duetector/tools/config_generator.py +++ b/duetector/tools/config_generator.py @@ -76,7 +76,7 @@ def __init__( for manager in self.managers: m = manager(include_extension=include_extension) _recursive_load(m.config_scope, self.dynamic_config, m.default_config) - for c in m.init(ignore_disabled=False): + for c in m.init(ignore_disabled=False, include_template=False): _recursive_load( c.config_scope, self.dynamic_config[m.config_scope], diff --git a/duetector/tracers/base.py b/duetector/tracers/base.py index c462448..189daec 100644 --- a/duetector/tracers/base.py +++ b/duetector/tracers/base.py @@ -240,7 +240,7 @@ class ShellTracer(Tracer): """ shell command """ - data_t = namedtuple("ShellOutput", ["output"]) + data_t = namedtuple("ShellOutput", ["output", "dt"]) """ data type for this tracer """ diff --git a/examples/README.md b/examples/README.md index ede2b6f..856ec19 100644 --- a/examples/README.md +++ b/examples/README.md @@ -3,3 +3,4 @@ ## Contents - [extension](./extension/): Duetector's extension examples +- [config](./config/): Examples of config files, including add tracer using config file. diff --git a/examples/config/sh-tracer-template.config.toml b/examples/config/sh-tracer-template.config.toml new file mode 100644 index 0000000..d0bc783 --- /dev/null +++ b/examples/config/sh-tracer-template.config.toml @@ -0,0 +1,9 @@ +# For more details of tracer template and tracer manager +# See: https://duetector.readthedocs.io/en/latest/managers/tracer.html + +[tracer.template] +disabled = false + +[tracer.template.sh] +# Add a pstracer to run the command "ps -aux" +pstracer = { "comm" = ["ps", "-aux"], config = { "enable_cache" = false } } diff --git a/tests/test_bcc_monitor.py b/tests/test_bcc_monitor.py index 6afa61a..523c712 100644 --- a/tests/test_bcc_monitor.py +++ b/tests/test_bcc_monitor.py @@ -114,8 +114,8 @@ def test_bcc_monitor(bcc_monitor: MockMonitor): bcc_monitor.poll_all() bcc_monitor.shutdown() assert bcc_monitor.summary() - bcc_monitor.summary()["MockMonitor"]["DBCollector"]["BccMockTracer"]["last"] == Tracking( - tracer="BccMockTracer", + bcc_monitor.summary()["MockMonitor"]["DBCollector"]["bccmocktracer"]["last"] == Tracking( + tracer="bccmocktracer", pid=9999, uid=9999, gid=9999, diff --git a/tests/test_tracer_template.py b/tests/test_tracer_template.py new file mode 100644 index 0000000..3f2fda5 --- /dev/null +++ b/tests/test_tracer_template.py @@ -0,0 +1,45 @@ +import pytest + +from duetector.managers.tracer import ShellTracer, TracerTemplate + + +def test_disabled_property(): + tracer_template = TracerTemplate({"disabled": True}) + assert tracer_template.disabled == True + + tracer_template = TracerTemplate({"disabled": False}) + assert tracer_template.disabled == False + + +def test_init(): + tracer_template = TracerTemplate({"disabled": True}) + assert tracer_template.init() == [] + + tracer_template = TracerTemplate( + { + "disabled": False, + "sh": { + "tracer_name": { + "comm": ["ps", "-aux"], + "config": {"enable_cache": False}, + } + }, + } + ) + tracers = tracer_template.init() + assert len(tracers) == 1 + assert isinstance(tracers[0], ShellTracer) + + +def test_get_wrap_tracer(): + tracer_template = TracerTemplate() + tracer = tracer_template._get_wrap_tracer( + "sh", "tracer_name", {"comm": ["ps", "-aux"], "config": {"enable_cache": False}} + ) + assert isinstance(tracer, ShellTracer) + assert tracer.config._config_dict == {"disabled": False, "enable_cache": False} + assert tracer.comm == ["ps", "-aux"] + + +if __name__ == "__main__": + pytest.main(["-s", "-vv", __file__])