tor-browser

The Tor Browser
git clone https://git.dasho.dev/tor-browser.git
Log | Files | Refs | README | LICENSE

websocketprocessbridge.py (3893B)

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