From c39c5453e9fd95a9ff6f9ec4be08c70cfd8241a3 Mon Sep 17 00:00:00 2001 From: Eric Leblond Date: Thu, 6 Dec 2012 11:21:57 +0100 Subject: [PATCH] suricatasc: refactor as a class The goal of this commit is to be able to use suricatasc has a library and and program. This is done by putting all active code in class and adding a Python magic to detect when file is used as a program. --- scripts/suricatasc/suricatasc.in | 344 ++++++++++++++++++------------- 1 file changed, 198 insertions(+), 146 deletions(-) diff --git a/scripts/suricatasc/suricatasc.in b/scripts/suricatasc/suricatasc.in index 31c6175936..a153f2d777 100755 --- a/scripts/suricatasc/suricatasc.in +++ b/scripts/suricatasc/suricatasc.in @@ -15,17 +15,44 @@ # Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA. import simplejson as json -import readline import re from socket import socket, AF_UNIX, error from time import sleep import sys -import argparse VERSION = "0.1" SIZE = 4096 -class Completer: +class SuricataException(Exception): + """ + Generic class for suricatasc exception + """ + def __init__(self, value): + self.value = value + + def __str__(self): + return str(self.value) + +class SuricataNetException(SuricataException): + """ + Exception raised when network error occur. + """ + pass + +class SuricataCommandException(SuricataException): + """ + Exception raised when command is not correct. + """ + pass + +class SuricataReturnException(SuricataException): + """ + Exception raised when return message is not correct. + """ + pass + + +class SuricataCompleter: def __init__(self, words): self.words = words self.generator = None @@ -44,155 +71,180 @@ class Completer: return None return None -def json_recv(socket): - cmdret = None - i = 0 - data = "" - while i < 5: - i += 1 - data += socket.recv(SIZE) +class SuricataSC: + def __init__(self, sck_path, verbose=False): + self.cmd_list=['shutdown','quit','pcap-file','pcap-file-number','pcap-file-list','iface-list','iface-stat'] + self.sck_path = sck_path + self.verbose = verbose + + def json_recv(self): + cmdret = None + i = 0 + data = "" + while i < 5: + i += 1 + data += self.socket.recv(SIZE) + try: + cmdret = json.loads(data) + break + except json.decoder.JSONDecodeError: + sleep(0.3) + return cmdret + + def send_command(self, command, arguments = None): + if command not in self.cmd_list and command != 'command-list': + raise SuricataCommandException("No such command: %s", command) + + cmdmsg = {} + cmdmsg['command'] = command + if (arguments != None): + cmdmsg['arguments'] = arguments + if self.verbose: + print "SND: " + json.dumps(cmdmsg) + self.socket.send(json.dumps(cmdmsg)) + cmdret = self.json_recv() + + if cmdret == None: + raise SuricataReturnException("Unable to get message from server") + + if self.verbose: + print "RCV: "+ json.dumps(cmdret) + + return cmdret + + def connect(self): try: - cmdret = json.loads(data) - break - except json.decoder.JSONDecodeError: - sleep(0.3) - return cmdret - -parser = argparse.ArgumentParser(prog='suricatasc', description='Client for Suricata unix socket') -parser.add_argument('-v', '--verbose', action='store_const', const=True, help='verbose output (including JSON dump)') -parser.add_argument('socket', metavar='socket', nargs='?', help='socket file to connnect to', default=None) -args = parser.parse_args() - -if args.socket != None: - SOCKET_PATH = "@e_localstatedir@/" + args.socket[0] -else: - SOCKET_PATH = "@e_localstatedir@/suricata-command.socket" - -socket = socket(AF_UNIX) -socket.connect(SOCKET_PATH) -socket.settimeout(10) - -#send version -if args.verbose: - print "SND: " + json.dumps({"version": VERSION}) -socket.send(json.dumps({"version": VERSION})) - -# get return -cmdret = json_recv(socket) - -if cmdret == None: - sys.stderr.write("Unable to get message from server") - sys.exit(1) + self.socket = socket(AF_UNIX) + self.socket.connect(SOCKET_PATH) + except error, err: + raise SuricataNetException(err) -if args.verbose: - print "RCV: "+ json.dumps(cmdret) -# if ok loop -if cmdret["return"] == "NOK": - sys.stderr.write("Error: %s" % (cmdret["message"])) - sys.exit(1) + self.socket.settimeout(10) + #send version + if self.verbose: + print "SND: " + json.dumps({"version": VERSION}) + self.socket.send(json.dumps({"version": VERSION})) -# get command list + # get return + cmdret = self.json_recv() -if args.verbose: - print "SND: " + json.dumps({"command": "command-list"}) -socket.send(json.dumps({"command": "command-list"})) -cmdret = json_recv(socket) + if cmdret == None: + raise SuricataReturnException("Unable to get message from server") -if cmdret == None: - sys.stderr.write("Unable to get message from server") - sys.exit(1) + if self.verbose: + print "RCV: "+ json.dumps(cmdret) -if args.verbose: - print "RCV: "+ json.dumps(cmdret) - - -if cmdret["return"] == "OK": - cmd_list = cmdret["message"]["commands"] - cmd_list.append("quit") - print "Command list: " + ", ".join(cmd_list) -else: - # This is the list of commands before command-list was added to the code. - cmd_list=['shutdown','quit','pcap-file','pcap-file-number','pcap-file-list','iface-list','iface-stat'] - -# if ok loop -try: - readline.set_completer(Completer(cmd_list)) - readline.set_completer_delims(";") - readline.parse_and_bind('tab: complete') - while True: - command = raw_input(">>> ").strip() - if command.split(' ', 2)[0] in cmd_list: - if command == "quit": - break; - cmdmsg = {} - if "pcap-file " in command: - try: - [cmd, filename, output] = command.split(' ', 2) - except: - print "Error: unable to split command '%s'" % (command) - continue - if cmd != "pcap-file": - print "Error: invalid command '%s'" % (command) - continue - else: - cmdmsg["command"] = cmd - cmdmsg["arguments"] = {} - cmdmsg["arguments"]["filename"] = filename - cmdmsg["arguments"]["output-dir"] = output - elif "iface-stat" in command: - try: - [cmd, iface] = command.split(' ', 1) - except: - print "Error: unable to split command '%s'" % (command) - continue - if cmd != "iface-stat": - print "Error: invalid command '%s'" % (command) - continue + if cmdret["return"] == "NOK": + raise SuricataReturnException("Error: %s" % (cmdret["message"])) + + def close(self): + self.socket.close() + + def interactive(self): + cmdret = self.send_command("command-list") + + # we silently ignore NOK as this means server is old + if cmdret["return"] == "OK": + self.cmd_list = cmdret["message"]["commands"] + self.cmd_list.append("quit") + print "Command list: " + ", ".join(self.cmd_list) + try: + readline.set_completer(SuricataCompleter(self.cmd_list)) + readline.set_completer_delims(";") + readline.parse_and_bind('tab: complete') + while True: + command = raw_input(">>> ").strip() + arguments = None + if command.split(' ', 2)[0] in self.cmd_list: + if command == "quit": + break; + if "pcap-file " in command: + try: + [cmd, filename, output] = command.split(' ', 2) + except: + print "Error: arguments to command '%s' is missing" % (command) + continue + if cmd != "pcap-file": + print "Error: invalid command '%s'" % (command) + continue + else: + arguments = {} + arguments["filename"] = filename + arguments["output-dir"] = output + elif "iface-stat" in command: + try: + [cmd, iface] = command.split(' ', 1) + except: + print "Error: unable to split command '%s'" % (command) + continue + if cmd != "iface-stat": + print "Error: invalid command '%s'" % (command) + continue + else: + arguments = {} + arguments["iface"] = iface + elif "conf-get" in command: + try: + [cmd, variable] = command.split(' ', 1) + except: + print "Error: unable to split command '%s'" % (command) + continue + if cmd != "conf-get": + print "Error: invalid command '%s'" % (command) + continue + else: + arguments = {} + arguments["variable"] = variable + else: + cmd = command else: - cmdmsg["command"] = cmd - cmdmsg["arguments"] = {} - cmdmsg["arguments"]["iface"] = iface - elif "conf-get" in command: - try: - [cmd, variable] = command.split(' ', 1) - except: - print "Error: unable to split command '%s'" % (command) - continue - if cmd != "conf-get": - print "Error: invalid command '%s'" % (command) + print "Error: unknown command '%s'" % (command) continue + + cmdret = self.send_command(cmd, arguments) + #decode json message + if cmdret["return"] == "NOK": + print "Error:" + print json.dumps(cmdret["message"], sort_keys=True, indent=4, separators=(',', ': ')) else: - cmdmsg["command"] = cmd - cmdmsg["arguments"] = {} - cmdmsg["arguments"]["variable"] = variable - else: - cmdmsg["command"] = command - if args.verbose: - print "SND: " + json.dumps(cmdmsg) - socket.send(json.dumps(cmdmsg)) - cmdret = json_recv(socket) - - if cmdret == None: - sys.stderr.write("Unable to get message from server") - sys.exit(1) - - if args.verbose: - print "RCV: "+ json.dumps(cmdret) - - #decode json message - if cmdret["return"] == "NOK": - print "Error:" - print json.dumps(cmdret["message"], sort_keys=True, indent=4, separators=(',', ': ')) - else: - print "Success:" - print json.dumps(cmdret["message"], sort_keys=True, indent=4, separators=(',', ': ')) - else: - print "Unknown command: '%s'" % (command) -except KeyboardInterrupt: - print "[!] Interrupted" - -print "[+] Quit command client" - -socket.close() - -sys.exit(1) + print "Success:" + print json.dumps(cmdret["message"], sort_keys=True, indent=4, separators=(',', ': ')) + except KeyboardInterrupt: + print "[!] Interrupted" + +if __name__ == '__main__': + import readline + import argparse + parser = argparse.ArgumentParser(prog='suricatasc', description='Client for Suricata unix socket') + parser.add_argument('-v', '--verbose', action='store_const', const=True, help='verbose output (including JSON dump)') + parser.add_argument('socket', metavar='socket', nargs='?', help='socket file to connnect to', default=None) + args = parser.parse_args() + + if args.socket != None: + SOCKET_PATH = "@e_localstatedir@/" + args.socket[0] + else: + SOCKET_PATH = "@e_localstatedir@/suricata-command.socket" + + sc = SuricataSC(SOCKET_PATH, verbose=args.verbose) + try: + sc.connect() + except SuricataNetException, err: + print "Unable to connect to socket %s: %s" % (SOCKET_PATH, err) + sys.exit(1) + except SuricataReturnException, err: + print "Unable to negotiate version with server: %s" % (err) + sys.exit(1) + try: + sc.interactive() + except SuricataNetException, err: + print "Communication error: %s" % (err) + sys.exit(1) + except SuricataReturnException, err: + print "Invalid return from server: %s" % (err) + sys.exit(1) + + print "[+] Quit command client" + + sc.close() + + sys.exit(1)