# The ZeekControl interactive shell. import logging import os import sys from ZeekControl import cmdresult, config, control, execute, lock, pluginreg, version from ZeekControl import node as node_mod from ZeekControl.exceptions import InvalidNodeError, LockError, RuntimeEnvironmentError class TermUI: def info(self, txt): print(txt) def error(self, txt): print(txt, file=sys.stderr) warn = error class NullHandler(logging.Handler): def emit(self, record): pass def expose(func): func.api_exposed = True return func def lock_required(func): def wrapper(self, *args, **kwargs): self.lock() try: return func(self, *args, **kwargs) finally: self.unlock() wrapper.lock_required = True return wrapper def lock_required_silent(func): def wrapper(self, *args, **kwargs): self.lock(showwait=False) try: return func(self, *args, **kwargs) finally: self.unlock() wrapper.lock_required = True return wrapper def check_config(func): def wrapper(self, *args, **kwargs): if config.Config.is_cfg_changed(): self.ui.warn('Configuration has changed. Run the "deploy" command.') return func(self, *args, **kwargs) return wrapper class ZeekCtl: def __init__( self, basedir=version.ZEEKBASE, libdir=version.LIBDIR, libdirinternal=version.LIBDIRINTERNAL, cfgfile=version.CFGFILE, zeekscriptdir=version.ZEEKSCRIPTDIR, ui=TermUI(), state=None, ): self.ui = ui self.zeekbase = basedir self.libdir = libdir self.libdirinternal = libdirinternal self.config = config.Configuration( self.zeekbase, self.libdir, self.libdirinternal, cfgfile, zeekscriptdir, self.ui, state, ) # Remove all log handlers (set by any previous calls to logging.*) logging.getLogger().handlers = [] if self.config.debug: # Add a log handler that logs to a file. try: logging.basicConfig( filename=self.config.debuglog, format="%(asctime)s [%(module)s] %(message)s", datefmt=self.config.timefmt, level=logging.DEBUG, ) except OSError as err: raise RuntimeEnvironmentError( f"{err}\nCheck if the user running ZeekControl has write access to the debug log file." ) else: # Add a log handler that does nothing. h = NullHandler() logging.getLogger().addHandler(h) self.executor = execute.Executor(self.config) self.plugins = pluginreg.PluginRegistry() self.setup() self.controller = control.Controller( self.config, self.ui, self.executor, self.plugins ) def setup(self): plugindirs = self.config.sitepluginpath.split(":") plugindirs.append(self.config.plugindir) plugindirs.append(self.config.pluginzeekdir) for pdir in plugindirs: if pdir: self.plugins.addDir(pdir) self.plugins.loadPlugins(self.ui, self.executor) self.plugins.initPluginOptions() self.plugins.addNodeKeys() self.config.initPostPlugins() self.plugins.initPlugins(self.ui) self.plugins.initPluginCmds() os.chdir(self.config.zeekbase) if self.config.get_state("cronenabled") is None: self.config.set_state("cronenabled", True) def reload_cfg(self): self.config.reload_cfg() if self.config.debug: if isinstance(logging.getLogger().handlers[0], NullHandler): # Remove the null handler and configure logging to a file. logging.getLogger().handlers = [] logging.basicConfig( filename=self.config.debuglog, format="%(asctime)s [%(module)s] %(message)s", datefmt=self.config.timefmt, level=logging.DEBUG, ) # Re-enable all log levels. logging.disable(logging.NOTSET) else: # Disable logging to all log levels that we use. logging.disable(logging.CRITICAL) self.executor.finish() self.plugins.initPluginOptions() self.config.initPostPlugins() self.plugins.initPlugins(self.ui) self.plugins.initPluginCmds() def finish(self): self.executor.finish() self.plugins.finishPlugins() def warn_zeekctl_install(self): self.config.warn_zeekctl_install() # Turns node name arguments into a list of nodes. If "get_hosts" is True, # then only one node per host is chosen. If "get_types" is True, then # only one node per node type (manager, proxy, etc.) is chosen. def node_args(self, args=None, get_hosts=False, get_types=False): nodes = [] if args: for arg in args.split(): nodelist = self.config.nodes(arg) if not nodelist: raise InvalidNodeError(f"unknown node '{arg}'") nodes += nodelist # Remove duplicate nodes newlist = list(set(nodes)) if len(newlist) != len(nodes): nodes = newlist else: # Get all nodes. nodes = self.config.nodes() # Sort the list so that it doesn't depend on initial order of arguments nodes.sort(key=node_mod.sortnode) if get_hosts: hosts = {} hostnodes = [] for node in nodes: if node.host not in hosts: hosts[node.host] = 1 hostnodes.append(node) nodes = hostnodes if get_types: types = {} typenodes = [] for node in nodes: if node.type not in types: types[node.type] = 1 typenodes.append(node) nodes = typenodes return nodes def lock(self, showwait=True): lockstatus = lock.lock(self.ui, showwait) if not lockstatus: raise LockError("Unable to get lock") self.config.read_state() def unlock(self): lock.unlock(self.ui) def node_names(self): return [n.name for n in self.config.nodes()] def node_groups(self): return node_mod.node_groups() @expose @check_config def nodes(self): results = cmdresult.CmdResult() if self.plugins.cmdPre("nodes"): for n in self.config.nodes(): results.set_node_data(n, True, n.to_dict()) else: results.ok = False self.plugins.cmdPost("nodes") return results @expose @check_config def get_config(self): results = cmdresult.CmdResult() if self.plugins.cmdPre("config"): results.keyval = self.config.options() else: results.ok = False self.plugins.cmdPost("config") return results @expose @check_config @lock_required def install(self, local=False): if self.plugins.cmdPre("install"): results = self.controller.install(local) else: results = cmdresult.CmdResult(ok=False) self.plugins.cmdPost("install") return results @expose @check_config @lock_required def start(self, node_list=None): nodes = self.node_args(node_list) nodes = self.plugins.cmdPreWithNodes("start", nodes) results = self.controller.start(nodes) self.plugins.cmdPostWithResults("start", results.get_node_data()) return results @expose @check_config @lock_required def stop(self, node_list=None): nodes = self.node_args(node_list) nodes = self.plugins.cmdPreWithNodes("stop", nodes) results = self.controller.stop(nodes) self.plugins.cmdPostWithResults("stop", results.get_node_data()) return results @expose @check_config @lock_required def restart(self, clean=False, node_list=None): nodes = self.node_args(node_list) nodes = self.plugins.cmdPreWithNodes("restart", nodes, clean) self.ui.info("stopping ...") results = self.stop(node_list) if not results.ok: return results if clean: self.ui.info("cleaning up ...") results = self.cleanup(node_list=node_list) if not results.ok: return results self.ui.info("checking configurations ...") results = self.check(node_list) if not results.ok: return results self.ui.info("installing ...") results = self.install() if not results.ok: return results self.ui.info("starting ...") results = self.start(node_list) self.plugins.cmdPostWithNodes("restart", nodes) return results @expose @lock_required def deploy(self): if not self.plugins.cmdPre("deploy"): results = cmdresult.CmdResult(ok=False) return results if self.config.is_cfg_changed(): self.ui.info("Reloading zeekctl configuration ...") self.reload_cfg() self.ui.info("checking configurations ...") results = self.check(check_node_types=True) if not results.ok: for node, success, output in results.get_node_output(): if not success: self.ui.info(f"{node} scripts failed.") self.ui.info(output) return results self.ui.info("installing ...") results = self.install() if not results.ok: return results self.ui.info("stopping ...") results = self.stop() if not results.ok: return results self.ui.info("starting ...") results = self.start() self.plugins.cmdPost("deploy") return results @expose @check_config @lock_required def status(self, node_list=None): nodes = self.node_args(node_list) nodes = self.plugins.cmdPreWithNodes("status", nodes) results = self.controller.status(nodes) self.plugins.cmdPostWithNodes("status", nodes) return results @expose @lock_required def top(self, node_list=None): nodes = self.node_args(node_list) nodes = self.plugins.cmdPreWithNodes("top", nodes) results = self.controller.top(nodes) self.plugins.cmdPostWithNodes("top", nodes) return results @expose @check_config @lock_required def diag(self, node_list=None): nodes = self.node_args(node_list) nodes = self.plugins.cmdPreWithNodes("diag", nodes) results = self.controller.diag(nodes) self.plugins.cmdPostWithNodes("diag", nodes) return results @expose @lock_required_silent def cron(self, watch=True): if self.plugins.cmdPre("cron", "", watch): self.controller.cron(watch) self.plugins.cmdPost("cron", "", watch) return True @expose @check_config @lock_required def cronenabled(self): results = False if self.plugins.cmdPre("cron", "?", False): results = self.config.cronenabled self.plugins.cmdPost("cron", "?", False) return results @expose @check_config @lock_required def setcronenabled(self, enable=True): if enable: if self.plugins.cmdPre("cron", "enable", False): self.config.set_state("cronenabled", True) self.ui.info("cron enabled") self.plugins.cmdPost("cron", "enable", False) else: if self.plugins.cmdPre("cron", "disable", False): self.config.set_state("cronenabled", False) self.ui.info("cron disabled") self.plugins.cmdPost("cron", "disable", False) return True @expose @check_config @lock_required def check(self, node_list=None, check_node_types=False): nodes = self.node_args(node_list, get_types=check_node_types) nodes = self.plugins.cmdPreWithNodes("check", nodes) results = self.controller.check(nodes) self.plugins.cmdPostWithResults("check", results.get_node_data()) return results @expose @check_config @lock_required def cleanup(self, cleantmp=False, node_list=None): nodes = self.node_args(node_list) nodes = self.plugins.cmdPreWithNodes("cleanup", nodes, cleantmp) results = self.controller.cleanup(nodes, cleantmp) self.plugins.cmdPostWithNodes("cleanup", nodes, cleantmp) return results @expose @check_config @lock_required def capstats(self, interval=10, node_list=None): nodes = self.node_args(node_list) nodes = self.plugins.cmdPreWithNodes("capstats", nodes, interval) results = self.controller.capstats(nodes, interval) self.plugins.cmdPostWithNodes("capstats", nodes, interval) return results @expose @check_config @lock_required def df(self, node_list=None): nodes = self.node_args(node_list, get_hosts=True) nodes = self.plugins.cmdPreWithNodes("df", nodes) results = self.controller.df(nodes) self.plugins.cmdPostWithNodes("df", nodes) return results @expose @check_config @lock_required def print_id(self, id, node_list=None): nodes = self.node_args(node_list) nodes = self.plugins.cmdPreWithNodes("print", nodes, id) results = self.controller.print_id(nodes, id) self.plugins.cmdPostWithNodes("print", nodes, id) return results @expose @check_config @lock_required def peerstatus(self, node_list=None): nodes = self.node_args(node_list) nodes = self.plugins.cmdPreWithNodes("peerstatus", nodes) results = self.controller.peerstatus(nodes) self.plugins.cmdPostWithNodes("peerstatus", nodes) return results @expose @check_config @lock_required def netstats(self, node_list=None): if not node_list: node_list = None if not self.config.standalone: node_list = node_mod.worker_group() nodes = self.node_args(node_list) nodes = self.plugins.cmdPreWithNodes("netstats", nodes) results = self.controller.netstats(nodes) self.plugins.cmdPostWithNodes("netstats", nodes) return results @expose @check_config def execute(self, cmd): nodes = self.node_args(get_hosts=True) if self.plugins.cmdPre("exec", cmd): results = self.controller.execute_cmd(nodes, cmd) else: results = cmdresult.CmdResult(ok=False) self.plugins.cmdPost("exec", cmd) return results @expose @check_config @lock_required def scripts(self, check=False, node_list=None): nodes = self.node_args(node_list) nodes = self.plugins.cmdPreWithNodes("scripts", nodes, check) results = self.controller.scripts(nodes, check) self.plugins.cmdPostWithNodes("scripts", nodes, check) return results @expose @check_config @lock_required def process(self, trace, options, scripts): if self.plugins.cmdPre("process", trace, options, scripts): results = self.controller.process(trace, options, scripts) else: results = cmdresult.CmdResult(ok=False) self.plugins.cmdPost("process", trace, options, scripts, results.ok) return results @expose @check_config @lock_required def plugincmd(self, cmd, args): return self.plugins.runCustomCommand(cmd, args, self.ui)