Newer
Older
eldiablo / diablo.py
GallaFrancesco on 10 Jun 4 KB comments + readme
#!/usr/bin/env python

# spawn process
from twisted.internet import protocol
from twisted.internet import reactor

# http web server
from twisted.web import server, resource
from twisted.internet import endpoints
from twisted.web.server import Site

# json encoding / decoding
import json

import sys
import logging

protoPID = dict()

class DispatchProtocol(protocol.ProcessProtocol):
    def __init__(self):
        self.stdout = b""
        self.stderr = b""
        self.ended = False
        self.exitcode = 0
        self.attached = False

    # process spawned
    def connectionMade(self):
        self.attached = True

    # stdout can be read from
    def outReceived(self, data):
        self.stdout = self.stdout + data

    # stderr can be read from
    def errReceived(self, data):
        self.stderr = self.stderr + data

    # process exited
    def processEnded(self, reason):
        self.exitcode = reason.value.exitCode
        self.ended = True

# TODO expand with these for better control
    def inConnectionLost(self):
        pass

    def outConnectionLost(self):
        pass

    def errConnectionLost(self):
        pass

    def processExited(self, reason):
        pass



def spawnTool(args): # TODO environmental variables
    protocol = DispatchProtocol() # initialize a process caller

    executable = args[0]
    process = reactor.spawnProcess(protocol, executable, args) # spawn executable

    protoPID[process.pid] = protocol # save the process id to get the protocol back
    return {
            "status": "LAUNCHED",
            "pid": process.pid
            }

def processStatus(pid):
    if not pid in protoPID: # invalid pid: error
        return {
                "status": "INVALID_PID",
                "pid": pid
                }
    protocol = protoPID[pid]
    if protocol.ended:
        if protocol.exitcode != 0: # process exited with error
            return {
                    "status": "EXIT_ERROR",
                    "pid": pid,
                    "code": protocol.exitcode,
                    "stdout": protocol.stdout.decode("utf-8"),
                    "stderr": protocol.stderr.decode("utf-8")
                    }
        else: # process exited cleanly
            return {
                    "status": "EXIT_SUCCESS",
                    "pid": pid,
                    "code": protocol.exitcode,
                    "stdout": protocol.stdout.decode("utf-8"),
                    "stderr": protocol.stderr.decode("utf-8")
                    }
    else: # process is still running
        return {
                "status": "RUNNING",
                "pid": pid,
                "stdout": protocol.stdout.decode("utf-8"),
                "stderr": protocol.stderr.decode("utf-8")
                }

class Commander(resource.Resource):
    isLeaf = True
    def render_POST(self, request):
        request.setHeader(b"Content-Type", b"<application/json")

        # decode the received command (should be { command: cmd, stdout: d })
        command = json.loads(request.content.read())
        ret = dict()

        if not "cmd" in command or not "data" in command: # invalid command received
            ret = { "status": "INVALID_COMMAND" }

        if command["cmd"] == "spawn": # spawn a tool on request
            args = command["data"].split("\n")
            ret = spawnTool(args)

        elif command["cmd"] == "status": # check process status
            if isinstance(command["data"], int) or command["data"].isdigit():
                ret = processStatus(int(command["data"]))

            else: # unable to convert pid to int
                ret = { "status": "INVALID_PID", "pid": command["data"] }

        else: # server error TODO handle
            ret = { "status": "DAEMON_ERROR" }

        # log received command TODO cleanup
        logging.info("[CMD] " +
                        command["cmd"] +
                        " " + str(command["data"]) +
                        " > pid: " +
                        str(ret["pid"]) +
                        ", status: " +
                        ret["status"])

        return json.dumps(ret).encode("utf-8")


if __name__ == "__main__":
    # default port
    port = 8080
    argv = sys.argv

    # informative log level
    logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s")

    # set custom port
    if len(argv) > 1:
        if argv[1] == "--port" or argv[1] == "-p":
            if not argv[2].isdigit() or int(argv[2]) > 65536:
                logging.error("[Error] Invalid port. Quitting.", file=sys.stderr)
                sys.exit(1)
            else:
                port = int(argv[2])

    # initialize the web server
    resource = Commander()
    site = Site(resource)
    endpoint = endpoints.TCP4ServerEndpoint(reactor, port)
    endpoint.listen(site)
    logging.info("[HTTP] Listening for requests on http://localhost:"+str(port))

    # listen for commands TODO daemonize
    reactor.run()