1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
|
# 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()
|
