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

@@ -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,13 +71,19 @@ class Completer:
             return None
         return None
 
-def json_recv(socket):
+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 += socket.recv(SIZE)
+            data += self.socket.recv(SIZE)
             try:
                 cmdret = json.loads(data)
                 break
@@ -58,87 +91,86 @@ def json_recv(socket):
                 sleep(0.3)
         return cmdret
 
-parser = argparse.ArgumentParser(prog='suricatasc', description='Client for Suricata unix socket')
+    def send_command(self, command, arguments = None):
-parser.add_argument('-v', '--verbose', action='store_const', const=True, help='verbose output (including JSON dump)')
+        if command not in self.cmd_list and command != 'command-list':
-parser.add_argument('socket', metavar='socket', nargs='?', help='socket file to connnect to', default=None)
+            raise SuricataCommandException("No such command: %s", command)
-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
+        cmdmsg = {}
-cmdret = json_recv(socket)
+        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:
+        if cmdret == None:
-    sys.stderr.write("Unable to get message from server")
+            raise SuricataReturnException("Unable to get message from server")
-    sys.exit(1)
 
-if args.verbose:
+        if self.verbose:
             print "RCV: "+ json.dumps(cmdret)
-# if ok loop
-if cmdret["return"] == "NOK":
-    sys.stderr.write("Error: %s" % (cmdret["message"]))
-    sys.exit(1)
 
-# get command list
+        return cmdret
 
-if args.verbose:
+    def connect(self):
-    print "SND: " + json.dumps({"command": "command-list"})
+        try:
-socket.send(json.dumps({"command": "command-list"}))
+            self.socket = socket(AF_UNIX)
-cmdret = json_recv(socket)
+            self.socket.connect(SOCKET_PATH)
+        except error, err:
+            raise SuricataNetException(err)
+
+        self.socket.settimeout(10)
+        #send version
+        if self.verbose:
+            print "SND: " + json.dumps({"version": VERSION})
+        self.socket.send(json.dumps({"version": VERSION}))
 
-if cmdret == None:
+        # get return
-    sys.stderr.write("Unable to get message from server")
+        cmdret = self.json_recv()
-    sys.exit(1)
+
+        if cmdret == None:
+            raise SuricataReturnException("Unable to get message from server")
 
-if args.verbose:
+        if self.verbose:
             print "RCV: "+ json.dumps(cmdret)
 
+        if cmdret["return"] == "NOK":
+            raise SuricataReturnException("Error: %s" % (cmdret["message"]))
+
+    def close(self):
+        self.socket.close()
 
-if cmdret["return"] == "OK":
+    def interactive(self):
-    cmd_list = cmdret["message"]["commands"]
+        cmdret = self.send_command("command-list")
-    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
+        # we silently ignore NOK as this means server is old
-try:
+        if cmdret["return"] == "OK":
-    readline.set_completer(Completer(cmd_list))
+            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()
-        if command.split(' ', 2)[0] in cmd_list:
+                arguments = None
+                if command.split(' ', 2)[0] in self.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)
+                            print "Error: arguments to command '%s' is missing" % (command)
                             continue
                         if cmd != "pcap-file":
                             print "Error: invalid command '%s'" % (command)
                             continue
                         else:
-                    cmdmsg["command"] = cmd
+                            arguments = {}
-                    cmdmsg["arguments"] = {}
+                            arguments["filename"] = filename
-                    cmdmsg["arguments"]["filename"] = filename
+                            arguments["output-dir"] = output
-                    cmdmsg["arguments"]["output-dir"] = output
                     elif "iface-stat" in command:
                         try:
                             [cmd, iface] = command.split(' ', 1)
@@ -149,9 +181,8 @@ try:
                             print "Error: invalid command '%s'" % (command)
                             continue
                         else:
-                    cmdmsg["command"] = cmd
+                            arguments = {}
-                    cmdmsg["arguments"] = {}
+                            arguments["iface"] = iface
-                    cmdmsg["arguments"]["iface"] = iface
                     elif "conf-get" in command:
                         try:
                             [cmd, variable] = command.split(' ', 1)
@@ -162,23 +193,15 @@ try:
                             print "Error: invalid command '%s'" % (command)
                             continue
                         else:
-                    cmdmsg["command"] = cmd
+                            arguments = {}
-                    cmdmsg["arguments"] = {}
+                            arguments["variable"] = variable
-                    cmdmsg["arguments"]["variable"] = variable
                     else:
-                cmdmsg["command"] = command
+                        cmd = command
-            if args.verbose:
+                else:
-                print "SND: " + json.dumps(cmdmsg)
+                    print "Error: unknown command '%s'" % (command)
-            socket.send(json.dumps(cmdmsg))
+                    continue
-            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)
 
+                cmdret = self.send_command(cmd, arguments)
                 #decode json message
                 if cmdret["return"] == "NOK":
                     print "Error:"
@@ -186,13 +209,42 @@ try:
                 else:
                     print "Success:"
                     print json.dumps(cmdret["message"], sort_keys=True, indent=4, separators=(',', ': '))
-        else:
+        except KeyboardInterrupt:
-            print "Unknown command: '%s'" % (command)
-except KeyboardInterrupt:
             print "[!] Interrupted"
 
-print "[+] Quit command client"
+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()
 
-socket.close()
+    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)
 
-sys.exit(1)
+    print "[+] Quit command client"
+
+    sc.close()
+
+    sys.exit(1)