#!/usr/bin/env python3 # Command-line interface for the Network Control Framework of Zeek, using Broker. import logging import netcontrol import sys import _thread from yaml import load try: from yaml import CLoader as Loader except ImportError: from yaml import Loader import string import re import argparse import sys from subprocess import check_output from subprocess import CalledProcessError #from future.utils import viewitems from enum import Enum, unique def parseArgs(): parser = argparse.ArgumentParser() parser.add_argument('--listen', default="127.0.0.1", help="Address to listen on for connections. Default: 127.0.0.1") parser.add_argument('--port', default=9977, help="Port to listen on for connections. Default: 9977") parser.add_argument('--topic', default="zeek/event/netcontrol-example", help="Topic to subscribe to. Default: zeek/event/netcontrol-example") parser.add_argument('--file', default="commands.yaml", help="File to read commands from. Default: commands.yaml") parser.add_argument('--debug', const=logging.DEBUG, default=logging.INFO, action='store_const', help="Enable debug output") args = parser.parse_args() return args class Listen: def __init__(self, queue, host, port, commands, **kwargs): self.logger = logging.getLogger(__name__) self.endpoint = netcontrol.Endpoint(queue, host, port) self.queuename = queue self.commands = commands self.use_threads = kwargs.get('use_threads', True) def listen_loop(self): self.logger.debug("Listen loop...") while 1==1: response = self.endpoint.getNextCommand() if response.type == netcontrol.ResponseType.AddRule: if self.use_threads: _thread.start_new_thread(self._add_remove_rule, (response, )) else: self._add_remove_rule(response) if response.type == netcontrol.ResponseType.RemoveRule: if self.use_threads: _thread.start_new_thread(self._add_remove_rule, (response, )) else: self._add_remove_rule(response) def _add_remove_rule(self, response): cmd = self.rule_to_cmd_dict(response.rule) if response.type == netcontrol.ResponseType.AddRule: type = 'add_rule' elif response.type == netcontrol.ResponseType.RemoveRule: type = 'remove_rule' else: self.logger.error("Internal error - incompabible rule type") type = 'unknown' if not ( type in self.commands): self.logger.error("No %s in commands", type) return commands = self.commands[type] output = "" self.logger.info("Received %s from Zeek: %s", type, cmd) for i in commands: currcmd = self.replace_command(i, cmd) output += "Command: "+currcmd+"\n" try: self.logger.info("Executing "+currcmd) cmdout = check_output(currcmd, shell=True) output += "Output: "+str(cmdout)+"\n" self.logger.debug("Command executed succefsully") except CalledProcessError as err: output = "Command "+currcmd+" failed with return code "+str(err.returncode)+" and output: "+str(err.output) self.logger.error(output) self.endpoint.sendRuleError(response, output) return except OSError as err: output = "Command "+currcmd+" failed with error code "+str(err.errno)+" ("+err.strerror+")" sel.logger.error(output) self.endpoint.sendRuleError(response, output) return if response.type == netcontrol.ResponseType.AddRule: self.logger.info("Sending rule_added to Zeek") self.endpoint.sendRuleAdded(response, output) else: self.logger.info("Sending rule_removed to Zeek") self.endpoint.sendRuleRemoved(response, output) def replace_single_command(self, argstr, cmds): reg = re.compile('\[(?P.)(?P.*?)(?:\:(?P.*?))?\]') #print argstr m = reg.search(argstr) if m == None: self.logger.error('%s could not be converted to rule', argstr) return '' type = m.group('type') target = m.group('target') arg = m.group('argument') if type == '?': if not ( target in cmds ): return '' elif arg == None: return cmds[target] # we have an argument *sigh* return re.sub(r'\.', cmds[target], arg) elif type == '!': if arg == None: self.logger.error("[!] needs argument for %s", argstr) return '' if not ( target in cmds ): return arg else: return '' else: self.logger.error("unknown command type %s in %s", type, argstr) return '' def replace_command(self, command, args): reg = re.compile('\[(?:\?|\!).*?\]') return reg.sub(lambda x: self.replace_single_command(x.group(), args), command) def rule_to_cmd_dict(self, rule): cmd = {} mapping = { 'type': 'ty', 'target': 'target', 'expire': 'expire', 'priority': 'priority', 'id': 'id', 'cid': 'cid', 'entity.ip': 'address', 'entity.mac': 'mac', 'entity.conn.orig_h': 'conn.orig_h', 'entity.conn.orig_p': 'conn.orig_p', 'entity.conn.resp_h': 'conn.resp_h', 'entity.conn.resp_p': 'conn.resp_p', 'entity.flow.src_h': 'flow.src_h', 'entity.flow.src_p': 'flow.src_p', 'entity.flow.dst_h': 'flow.dst_h', 'entity.flow.dst_p': 'flow.dst_p', 'entity.flow.src_m': 'flow.src_m', 'entity.flow.dst_m': 'flow.dst_m', 'entity.mod.src_h': 'mod.src_h', 'entity.mod.src_p': 'mod.src_p', 'entity.mod.dst_h': 'mod.dst_h', 'entity.mod.dst_p': 'mod.dst_p', 'entity.mod.src_m': 'mod.src_m', 'entity.mod.dst_m': 'mod.dst_m', 'entity.mod.redirect_port': 'mod.port', 'entity.i': 'mod.port', } for (k, v) in list(mapping.items()): path = k.split('.') e = rule for i in path: if e == None: break elif i in e: e = e[i] else: e = None break if e == None: continue if isinstance(e, tuple): cmd[v] = e[0] cmd[v+".proto"] = e[1] else: cmd[v] = e if isinstance(e, str): spl = e.split("/") if len(spl) > 1: cmd[v+".ip"] = spl[0] cmd[v+".net"] = spl[1] proto = mapping.get('entity.conn.orig_p.proto', mapping.get('entity.conn.dest_p.proto', mapping.get('entity.flow.src_p.proto', mapping.get('entity.flow.dst_p.proto', None)))) if proto != None: entity['proto'] = proto return cmd args = parseArgs() stream = open(args.file, 'r') config = load(stream, Loader=Loader) logging.basicConfig(level=args.debug) logging.info("Starting command-line client...") zeekcon = Listen(args.topic, args.listen, int(args.port), config) zeekcon.listen_loop()