Feature: template for tracer (#76)

* Support include_extension and init template

* Suuport ShTracer template

* Add example for tracer template

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`

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:

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

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": {},

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.
 

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

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.
 

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.
 

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,13 +23,94 @@ 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.
 
     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

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 _
 

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"

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],

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
     """

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.

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 } }

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,

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__])