File: pluginmanager.py

package info (click to toggle)
prelude-correlator 4.1.1-2
  • links: PTS, VCS
  • area: main
  • in suites: buster
  • size: 1,056 kB
  • sloc: python: 1,686; sh: 34; makefile: 23
file content (214 lines) | stat: -rw-r--r-- 7,097 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
# Copyright (C) 2009-2017 CS-SI. All Rights Reserved.
# Author: Yoann Vandoorselaere <yoann.v@prelude-ids.com>
#
# This file is part of the Prelude-Correlator program.
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2, or (at your option)
# any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.

import pkg_resources
import os
import imp

from preludecorrelator import log, error, require, plugins

logger = log.getLogger(__name__)


class Plugin(object):
    enable = True
    autoload = True
    conflict = []

    def getConfigValue(self, option, fallback=None, type=str):
        return self.env.config.get(self.__class__.__name__, option, fallback=fallback, type=type)

    def __init__(self, env):
        self.env = env

    def _getName(self):
        return self.__class__.__name__

    def stats(self):
        pass

    def signal(self, signo, frame):
        pass

    def run(self, idmef):
        pass


class PluginDependenciesError(ImportError):
     pass

class PluginManager(object):
    _default_entrypoint = 'preludecorrelator.plugins'

    def __init__(self, env, entrypoint=None):
        self._env = env
        self._count = 0
        self.__plugins_instances = []
        self.__plugins_classes = []

        self._conflict = {}
        self._force_enable = {}

        plugin_entries =  [ (e.name, e, self._load_entrypoint) for e in pkg_resources.iter_entry_points(entrypoint if entrypoint else self._default_entrypoint)]
        if entrypoint is None:
            plugin_entries += [ (u[0], u, self._load_userpoint ) for u in self._get_userpoints(env)]

        for pname, e, fct in plugin_entries:
            logger.debug("loading point %s", pname, level=1)

            enable_s = env.config.get(pname, "enable", fallback=None)
            if enable_s:
                enable_s = enable_s.lower()

            enable = enable_s in ("true", "yes", "force", None)
            disable = env.config.getAsBool(pname, "disable", fallback=False)

            # do not load if the user specifically used disable=true, or enable=false
            if not enable or disable:
                logger.info("[%s]: disabled on user request", pname)
                continue

            plugin_class = fct(e)
            
            if plugin_class is None:
                continue

            if not enable_s:
                enable = plugin_class.enable

            if enable:
                if disable:
                    enable = False

                elif enable_s == "force":
                    self._force_enable[pname] = enable

            if not enable:
                logger.info("[%s]: disabled by default", pname)
                continue

            for reason, namelist in plugin_class.conflict:
                self._conflict.update([(name, (pname, reason)) for name in namelist])

            self.__plugins_classes.append(plugin_class)

    def load(self):
        for plugin_class in self.getPluginsClassesList():
            pname = plugin_class.__name__

            if pname in self._conflict and not pname in self._force_enable:
                logger.info("[%s]: disabled by plugin '%s' reason:%s", pname, self._conflict[pname][0], self._conflict[pname][1])
                continue

            if plugin_class.autoload:
                try:
                    pi = plugin_class(self._env)

                except error.UserError as e:
                    logger.error("[%s]: %s", pname, e)
                    raise error.UserError("Plugin '%s' failed to load, please fix the issue or disable the plugin" % pname)

                self.__plugins_instances.append(pi)

            self._count += 1

    def _get_userpoints(self, env):
        if not env.config.has_section("python_rules"):
            python_rules_dirs = require.get_config_filename("rules/python")
        else:
            python_rules_dirs = env.config.get("python_rules", "paths", fallback="")

        for pathdir in python_rules_dirs.splitlines():
            if not os.access(pathdir, os.R_OK) or not os.path.isdir(pathdir):
                logger.warning("Can not load %s python rules dir" % pathdir)
                continue

            for f in os.listdir(pathdir):
                if not f.endswith('.py') or f == '__init__.py':
                    continue

                if os.path.isdir(os.path.join(pathdir, f)):
                    continue

                yield (f.rpartition('.')[0], pathdir)

    def _load_entrypoint(self, entrypoint):
        try:
            return entrypoint.load()

        except ImportError as e:
            logger.error("[%s]: import error: %s", entrypoint.name, e)
            return None

        except Exception as e:
            logger.exception("[%s]: loading error: %s", entrypoint.name, e)
            return None

    def _load_userpoint(self, args):
        name, path = args
        try:
            mod_info = imp.find_module(name, [path])

        except ImportError:
            logger.warning( 'Invalid plugin "%s" in "%s"' % (name, path) )
            return None

        try:
            return getattr(imp.load_module( self._default_entrypoint + '.' + name , *mod_info), name)

        except Exception as e:
            logger.warning( "Unable to load %(file)s: %(error)s" % {'file': name,'error': str(e),})
            return None

    def getPluginCount(self):
        return self._count

    def getPluginList(self):
        return self.getPluginsInstancesList()

    def getPluginsInstancesList(self):
        return self.__plugins_instances

    def getPluginsClassesList(self):
        return self.__plugins_classes

    def stats(self):
        for plugin in self.getPluginsInstancesList():
            try:
                plugin.stats()
            except Exception:
                logger.exception("[%s]: exception occurred while retrieving statistics", plugin._getName())

    def signal(self, signo, frame):
        for plugin in self.getPluginsInstancesList():
            try:
                plugin.signal(signo, frame)
            except Exception:
                logger.exception("[%s]: exception occurred while signaling", plugin._getName())

    def run(self, idmef):
        for plugin in self.getPluginsInstancesList():
            try:
                plugin.run(idmef)

            except error.UserError as e:
                logger.error("[%s]: error running plugin : %s", plugin._getName(), e)

            except Exception:
                logger.exception("[%s]: exception occurred while running", plugin._getName())