123 lines
3.8 KiB
Python
123 lines
3.8 KiB
Python
# vim: set ts=4 et sw=4 tw=80
|
|
# This Source Code Form is subject to the terms of the Mozilla Public
|
|
# License, v. 2.0. If a copy of the MPL was not distributed with this
|
|
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
|
|
|
|
from twisted.internet import protocol, reactor
|
|
from twisted.internet.task import LoopingCall
|
|
from autobahn.twisted.websocket import WebSocketServerProtocol, WebSocketServerFactory
|
|
|
|
import psutil
|
|
|
|
import argparse
|
|
import six
|
|
import sys
|
|
import os
|
|
|
|
# maps a command issued via websocket to running an executable with args
|
|
commands = {
|
|
"iceserver": [sys.executable, "-u", os.path.join("iceserver", "iceserver.py")]
|
|
}
|
|
|
|
|
|
class ProcessSide(protocol.ProcessProtocol):
|
|
"""Handles the spawned process (I/O, process termination)"""
|
|
|
|
def __init__(self, socketSide):
|
|
self.socketSide = socketSide
|
|
|
|
def outReceived(self, data):
|
|
data = six.ensure_str(data)
|
|
if self.socketSide:
|
|
lines = data.splitlines()
|
|
for line in lines:
|
|
self.socketSide.sendMessage(line.encode("utf8"), False)
|
|
|
|
def errReceived(self, data):
|
|
self.outReceived(data)
|
|
|
|
def processEnded(self, reason):
|
|
if self.socketSide:
|
|
self.outReceived(reason.getTraceback())
|
|
self.socketSide.processGone()
|
|
|
|
def socketGone(self):
|
|
self.socketSide = None
|
|
self.transport.loseConnection()
|
|
self.transport.signalProcess("KILL")
|
|
|
|
|
|
class SocketSide(WebSocketServerProtocol):
|
|
"""
|
|
Handles the websocket (I/O, closed connection), and spawning the process
|
|
"""
|
|
|
|
def __init__(self):
|
|
super(SocketSide, self).__init__()
|
|
self.processSide = None
|
|
|
|
def onConnect(self, request):
|
|
return None
|
|
|
|
def onOpen(self):
|
|
return None
|
|
|
|
def onMessage(self, payload, isBinary):
|
|
# We only expect a single message, which tells us what kind of process
|
|
# we're supposed to launch. ProcessSide pipes output to us for sending
|
|
# back to the websocket client.
|
|
if not self.processSide:
|
|
self.processSide = ProcessSide(self)
|
|
# We deliberately crash if |data| isn't on the "menu",
|
|
# or there is some problem spawning.
|
|
data = six.ensure_str(payload)
|
|
try:
|
|
reactor.spawnProcess(
|
|
self.processSide, commands[data][0], commands[data], env=os.environ
|
|
)
|
|
except BaseException as e:
|
|
print(e.str())
|
|
self.sendMessage(e.str())
|
|
self.processGone()
|
|
|
|
def onClose(self, wasClean, code, reason):
|
|
if self.processSide:
|
|
self.processSide.socketGone()
|
|
|
|
def processGone(self):
|
|
self.processSide = None
|
|
self.transport.loseConnection()
|
|
|
|
|
|
# Parent process could have already exited, so this is slightly racy. Only
|
|
# alternative is to set up a pipe between parent and child, but that requires
|
|
# special cooperation from the parent.
|
|
parent_process = psutil.Process(os.getpid()).parent()
|
|
|
|
|
|
def check_parent():
|
|
"""Checks if parent process is still alive, and exits if not"""
|
|
if not parent_process.is_running():
|
|
print("websocket/process bridge exiting because parent process is gone")
|
|
reactor.stop()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Starts websocket/process bridge.")
|
|
parser.add_argument(
|
|
"--port",
|
|
type=str,
|
|
dest="port",
|
|
default="8191",
|
|
help="Port for websocket/process bridge. Default 8191.",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
parent_checker = LoopingCall(check_parent)
|
|
parent_checker.start(1)
|
|
|
|
bridgeFactory = WebSocketServerFactory()
|
|
bridgeFactory.protocol = SocketSide
|
|
reactor.listenTCP(int(args.port), bridgeFactory)
|
|
print("websocket/process bridge listening on port %s" % args.port)
|
|
reactor.run()
|