snowmix/pmanager/pmanager.py

#!/usr/bin/env python3
import json
import os
import signal
import socket
import subprocess
import sys
import time
import websocket

DEBUG = True
DEBUG = False
def log(s):
    if not DEBUG:
        return
    print(s)

def sendmsg(msg):
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    ret = None
    try:
        s.connect(('localhost', 9999))
        tosend = msg + "\r\n"
        encoded_msg = tosend.encode("ascii")
        s.sendall(encoded_msg)
        ret = s.recv(4096)
        s.close()
    except Exception as e:
        log(e)
        return None
    log("[nc output] %s " % ret)
    return ret

class PManager:
    ws = websocket.WebSocket()
    ip = "ws://192.168.1.208:9001"
    p_list = {}
    ptable = {}
    log_list = {}
    name = ""
    def __init__(self, name, ptable):
        self.ptable = ptable
        self.conn(self.ip)
        self.name = name

        signal.signal(signal.SIGINT, self.sighandler)
        signal.signal(signal.SIGCHLD, signal.SIG_IGN) #To kill zombies
        while True:
            try:
                data = self.ws.recv()
            except Exception as e:
                log(e)
                self.conn(self.ip)
                continue
            try:
                j = json.loads(data)
            except Exception as e:
                log(e)
                log("Invalid json?")
                continue
            
            if ("type" not in j or "command" not in j or
                    "key" not in j or j["type"] != "snowmix"):
                log("Dropping: %s" % j)
                continue
        
            #log(j)
            if j["command"] == "oneshot":
                self.oneshot(j["key"])

            if j["command"] == "launch":
                opts = None
                if "options" in j:
                    opts = j["options"]
                self.launch(j["key"],opts)
        
            if j["command"] == "kill":
                self.stop(j["key"])
            
            if j["command"] == "netcat":
                sendmsg(j["key"])
            if j["command"] == "status": #and j["key"] == self.name:
                self.send({"type": "snowmix", "type": "statusreport", "name": self.name, "data":self.p_list})

    def send(self, data):
        try:
            if type(data) is dict:
                data = json.dumps(data)
            print(data)
            self.ws.send(data)
        except Exception as e:
            print(e)

    def conn(self, target):
        if self.ws is None:
            self.ws = websocket.WebSocket()
        try:
            self.ws.connect(target)
        except Exception as e:
            log(e)
            time.sleep(2)
            self.conn(target) #keep trying you basterd
    
    
    def close(self):
        if self.ws is None:
            return
        self.ws.close()
    
    def sighandler(self, sig, frame):
        log("\rExiting!")
        self.close()
        log("Bye!")
        sys.exit(0)
    
    def stop(self, key):
        if key not in self.p_list:
            log("key %s doesn't exist! not killing!" % (key))
            return
        try:
            os.killpg(os.getpgid(self.p_list[key]), signal.SIGTERM)
        except Exception as e:
            log("This process (%d) can't be killed! Why?" % self.p_list[key])
            log(e)
    
        try:
            self.log_list[key].close()
        except Exception as e:
            log("ex file: %s" % e)
    
        if key in self.log_list:
            del self.log_list[key]
    
        del self.p_list[key]
            
    def oneshot(self, key):
        if key not in self.ptable:
            log("This process (%s) is not on my table!" % key)
            return
        log("Oneshotting! key: %s" % key)
        try:
            proc = subprocess.Popen(self.ptable[key], preexec_fn=os.setsid)
        except Exception as e:
            log(e)

    def launch(self, key, options=None):
        if key in self.p_list:
            log("key %s exist! not running!" % (key))
            return
        
        if key not in self.ptable:
            log("This process (%s) is not on my table!" % key)
            return
    
        log("Launching! key: %s, opts: %s" % (key, options))
        tolaunch = self.ptable[key]
        if options is not None:
            options = [ o for o in options if type(o) is str ]
            tolaunch = tolaunch + options
    
        log(tolaunch)
        self.log_list[key] = open("/tmp/%s" % key, "w")
        try:
            proc = subprocess.Popen(tolaunch, preexec_fn=os.setsid, stdout=self.log_list[key])
            #Le doy su propio grupo al subproceso
        except Exception as e:
            log(e)
            sys.exit(1)

        self.p_list[key] = proc.pid
        return proc.pid

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: ./%s <file.json>" % (sys.argv[0]))
        sys.exit(1)
    try:
        data = json.loads(open(sys.argv[1]).read())
    except Exception as e:
        print(e)
        sys.exit(2)
    
    ptable = data["data"]
    name = data["name"]
    m = PManager(name, ptable)