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