1# vim: set ts=4 et sw=4 tw=80 2# This Source Code Form is subject to the terms of the Mozilla Public 3# License, v. 2.0. If a copy of the MPL was not distributed with this 4# file, You can obtain one at http://mozilla.org/MPL/2.0/. 5 6from __future__ import absolute_import, print_function 7 8from twisted.internet import protocol, reactor 9from twisted.internet.task import LoopingCall 10from autobahn.twisted.websocket import WebSocketServerProtocol, WebSocketServerFactory 11 12import psutil 13 14import argparse 15import six 16import sys 17import os 18 19# maps a command issued via websocket to running an executable with args 20commands = { 21 "iceserver": [sys.executable, "-u", os.path.join("iceserver", "iceserver.py")] 22} 23 24 25class ProcessSide(protocol.ProcessProtocol): 26 """Handles the spawned process (I/O, process termination)""" 27 28 def __init__(self, socketSide): 29 self.socketSide = socketSide 30 31 def outReceived(self, data): 32 data = six.ensure_str(data) 33 if self.socketSide: 34 lines = data.splitlines() 35 for line in lines: 36 self.socketSide.sendMessage(line.encode("utf8"), False) 37 38 def errReceived(self, data): 39 self.outReceived(data) 40 41 def processEnded(self, reason): 42 if self.socketSide: 43 self.outReceived(reason.getTraceback()) 44 self.socketSide.processGone() 45 46 def socketGone(self): 47 self.socketSide = None 48 self.transport.loseConnection() 49 self.transport.signalProcess("KILL") 50 51 52class SocketSide(WebSocketServerProtocol): 53 """ 54 Handles the websocket (I/O, closed connection), and spawning the process 55 """ 56 57 def __init__(self): 58 super(SocketSide, self).__init__() 59 self.processSide = None 60 61 def onConnect(self, request): 62 return None 63 64 def onOpen(self): 65 return None 66 67 def onMessage(self, payload, isBinary): 68 # We only expect a single message, which tells us what kind of process 69 # we're supposed to launch. ProcessSide pipes output to us for sending 70 # back to the websocket client. 71 if not self.processSide: 72 self.processSide = ProcessSide(self) 73 # We deliberately crash if |data| isn't on the "menu", 74 # or there is some problem spawning. 75 data = six.ensure_str(payload) 76 try: 77 reactor.spawnProcess( 78 self.processSide, commands[data][0], commands[data], env=os.environ 79 ) 80 except BaseException as e: 81 print(e.str()) 82 self.sendMessage(e.str()) 83 self.processGone() 84 85 def onClose(self, wasClean, code, reason): 86 if self.processSide: 87 self.processSide.socketGone() 88 89 def processGone(self): 90 self.processSide = None 91 self.transport.loseConnection() 92 93 94# Parent process could have already exited, so this is slightly racy. Only 95# alternative is to set up a pipe between parent and child, but that requires 96# special cooperation from the parent. 97parent_process = psutil.Process(os.getpid()).parent() 98 99 100def check_parent(): 101 """ Checks if parent process is still alive, and exits if not """ 102 if not parent_process.is_running(): 103 print("websocket/process bridge exiting because parent process is gone") 104 reactor.stop() 105 106 107if __name__ == "__main__": 108 parser = argparse.ArgumentParser(description="Starts websocket/process bridge.") 109 parser.add_argument( 110 "--port", 111 type=str, 112 dest="port", 113 default="8191", 114 help="Port for websocket/process bridge. Default 8191.", 115 ) 116 args = parser.parse_args() 117 118 parent_checker = LoopingCall(check_parent) 119 parent_checker.start(1) 120 121 bridgeFactory = WebSocketServerFactory() 122 bridgeFactory.protocol = SocketSide 123 reactor.listenTCP(int(args.port), bridgeFactory) 124 print("websocket/process bridge listening on port %s" % args.port) 125 reactor.run() 126