suricata/scripts/suricatasc/suricatasc.in

#!/usr/bin/python
# Copyright(C) 2012 Open Information Security Foundation

# 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, version 2 of the License.
#
# 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., 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

class Completer:
    def __init__(self, words):
        self.words = words
        self.generator = None

    def complete(self, text):
        for word in self.words:
            if word.startswith(text):
                yield word

    def __call__(self, text, state):
        if state == 0:
            self.generator = self.complete(text)
        try:
            return self.generator.next()
        except StopIteration:
            return None
        return None

def json_recv(socket):
    cmdret = None
    i = 0
    data = ""
    while i < 5:
        i += 1
        data += socket.recv(SIZE)
        try:
            cmdret = json.loads(data)
            break
        except json.decoder.JSONDecodeError:
            sleep(0.3)
    return cmdret

VERSION = "0.1"
if len(sys.argv) == 2:
    SOCKET_PATH = "@e_localstatedir@/" + sys.argv[1]
else:
    SOCKET_PATH = "@e_localstatedir@/suricata-command.socket"
SIZE = 4096

socket = socket(AF_UNIX)
socket.connect(SOCKET_PATH)
socket.settimeout(10)

#send 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)

# if ok loop
if cmdret["return"] == "NOK":
    sys.stderr.write("Error: %s" % (cmdret["message"]))
    sys.exit(1)

# get command list

socket.send(json.dumps({"command": "command-list"}))
cmdret = json_recv(socket)

if cmdret == None:
    sys.stderr.write("Unable to get message from server")
    sys.exit(1)

# if ok loop
if cmdret["return"] == "NOK":
    sys.stderr.write("Error: %s" % (cmdret["message"]))
    sys.exit(1)

cmd_list = cmdret["message"]["commands"]
cmd_list.append("quit")

# 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
                else:
                    cmdmsg["command"] = cmd
                    cmdmsg["arguments"] = {}
                    cmdmsg["arguments"]["iface"] = iface
            else:
                cmdmsg["command"] = command
            socket.send(json.dumps(cmdmsg))
            cmdret = json_recv(socket)

            if cmdret == None:
                sys.stderr.write("Unable to get message from server")
                sys.exit(1)
            #decode json message
            if cmdret["return"] == "NOK":
                print "Error: %s" % (cmdret["message"])
            else:
                print "Success: %s" % (cmdret["message"])
        else:
            print "Unknown command: '%s'" % (command)
except KeyboardInterrupt:
    print "[!] Interrupted"

print "[+] Quit command client"

socket.close()

sys.exit(1)