← Back to branch summary

~ubuntu-branches/ubuntu/vivid/emscripten/vivid

« back to all changes in this revision

Viewing changes to third_party/websockify/websockify/websocketproxy.py

  • Committer: Package Import Robot
  • Author(s): Sylvestre Ledru
  • Date: 2013-05-02 13:11:51 UTC
  • Revision ID: package-import@ubuntu.com-20130502131151-q8dvteqr1ef2x7xz
Tags: upstream-1.4.1~20130504~adb56cb
Import upstream version 1.4.1~20130504~adb56cb

Show diffs side-by-side

added added

removed removed

Lines of Context:
third_party/websockify/websockify/websocketproxy.py
 
1
#!/usr/bin/env python2
 
2
 
 
3
'''
 
4
A WebSocket to TCP socket proxy with support for "wss://" encryption.
 
5
Copyright 2011 Joel Martin
 
6
Licensed under LGPL version 3 (see docs/LICENSE.LGPL-3)
 
7
 
 
8
You can make a cert/key with openssl using:
 
9
openssl req -new -x509 -days 365 -nodes -out self.pem -keyout self.pem
 
10
as taken from http://docs.python.org/dev/library/ssl.html#certificates
 
11
 
 
12
'''
 
13
 
 
14
import signal, socket, optparse, time, os, sys, subprocess
 
15
from select import select
 
16
import websocket
 
17
try:
 
18
    from urllib.parse import parse_qs, urlparse
 
19
except:
 
20
    from cgi import parse_qs
 
21
    from urlparse import urlparse
 
22
 
 
23
class WebSocketProxy(websocket.WebSocketServer):
 
24
    """
 
25
    Proxy traffic to and from a WebSockets client to a normal TCP
 
26
    socket server target. All traffic to/from the client is base64
 
27
    encoded/decoded to allow binary data to be sent/received to/from
 
28
    the target.
 
29
    """
 
30
 
 
31
    buffer_size = 65536
 
32
 
 
33
    traffic_legend = """
 
34
Traffic Legend:
 
35
    }  - Client receive
 
36
    }. - Client receive partial
 
37
    {  - Target receive
 
38
 
 
39
    >  - Target send
 
40
    >. - Target send partial
 
41
    <  - Client send
 
42
    <. - Client send partial
 
43
"""
 
44
 
 
45
    def __init__(self, *args, **kwargs):
 
46
        # Save off proxy specific options
 
47
        self.target_host    = kwargs.pop('target_host', None)
 
48
        self.target_port    = kwargs.pop('target_port', None)
 
49
        self.wrap_cmd       = kwargs.pop('wrap_cmd', None)
 
50
        self.wrap_mode      = kwargs.pop('wrap_mode', None)
 
51
        self.unix_target    = kwargs.pop('unix_target', None)
 
52
        self.ssl_target     = kwargs.pop('ssl_target', None)
 
53
        self.target_cfg     = kwargs.pop('target_cfg', None)
 
54
        # Last 3 timestamps command was run
 
55
        self.wrap_times    = [0, 0, 0]
 
56
 
 
57
        if self.wrap_cmd:
 
58
            rebinder_path = ['./', os.path.dirname(sys.argv[0])]
 
59
            self.rebinder = None
 
60
 
 
61
            for rdir in rebinder_path:
 
62
                rpath = os.path.join(rdir, "rebind.so")
 
63
                if os.path.exists(rpath):
 
64
                    self.rebinder = rpath
 
65
                    break
 
66
 
 
67
            if not self.rebinder:
 
68
                raise Exception("rebind.so not found, perhaps you need to run make")
 
69
            self.rebinder = os.path.abspath(self.rebinder)
 
70
 
 
71
            self.target_host = "127.0.0.1"  # Loopback
 
72
            # Find a free high port
 
73
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 
74
            sock.bind(('', 0))
 
75
            self.target_port = sock.getsockname()[1]
 
76
            sock.close()
 
77
 
 
78
            os.environ.update({
 
79
                "LD_PRELOAD": self.rebinder,
 
80
                "REBIND_OLD_PORT": str(kwargs['listen_port']),
 
81
                "REBIND_NEW_PORT": str(self.target_port)})
 
82
 
 
83
        if self.target_cfg:
 
84
            self.target_cfg = os.path.abspath(self.target_cfg)
 
85
 
 
86
        websocket.WebSocketServer.__init__(self, *args, **kwargs)
 
87
 
 
88
    def run_wrap_cmd(self):
 
89
        print("Starting '%s'" % " ".join(self.wrap_cmd))
 
90
        self.wrap_times.append(time.time())
 
91
        self.wrap_times.pop(0)
 
92
        self.cmd = subprocess.Popen(
 
93
                self.wrap_cmd, env=os.environ, preexec_fn=_subprocess_setup)
 
94
        self.spawn_message = True
 
95
 
 
96
    def started(self):
 
97
        """
 
98
        Called after Websockets server startup (i.e. after daemonize)
 
99
        """
 
100
        # Need to call wrapped command after daemonization so we can
 
101
        # know when the wrapped command exits
 
102
        if self.wrap_cmd:
 
103
            dst_string = "'%s' (port %s)" % (" ".join(self.wrap_cmd), self.target_port)
 
104
        elif self.unix_target:
 
105
            dst_string = self.unix_target
 
106
        else:
 
107
            dst_string = "%s:%s" % (self.target_host, self.target_port)
 
108
 
 
109
        if self.target_cfg:
 
110
            msg = "  - proxying from %s:%s to targets in %s" % (
 
111
                self.listen_host, self.listen_port, self.target_cfg)
 
112
        else:
 
113
            msg = "  - proxying from %s:%s to %s" % (
 
114
                self.listen_host, self.listen_port, dst_string)
 
115
 
 
116
        if self.ssl_target:
 
117
            msg += " (using SSL)"
 
118
 
 
119
        print(msg + "\n")
 
120
 
 
121
        if self.wrap_cmd:
 
122
            self.run_wrap_cmd()
 
123
 
 
124
    def poll(self):
 
125
        # If we are wrapping a command, check it's status
 
126
 
 
127
        if self.wrap_cmd and self.cmd:
 
128
            ret = self.cmd.poll()
 
129
            if ret != None:
 
130
                self.vmsg("Wrapped command exited (or daemon). Returned %s" % ret)
 
131
                self.cmd = None
 
132
 
 
133
        if self.wrap_cmd and self.cmd == None:
 
134
            # Response to wrapped command being gone
 
135
            if self.wrap_mode == "ignore":
 
136
                pass
 
137
            elif self.wrap_mode == "exit":
 
138
                sys.exit(ret)
 
139
            elif self.wrap_mode == "respawn":
 
140
                now = time.time()
 
141
                avg = sum(self.wrap_times)/len(self.wrap_times)
 
142
                if (now - avg) < 10:
 
143
                    # 3 times in the last 10 seconds
 
144
                    if self.spawn_message:
 
145
                        print("Command respawning too fast")
 
146
                        self.spawn_message = False
 
147
                else:
 
148
                    self.run_wrap_cmd()
 
149
 
 
150
    #
 
151
    # Routines above this point are run in the master listener
 
152
    # process.
 
153
    #
 
154
 
 
155
    #
 
156
    # Routines below this point are connection handler routines and
 
157
    # will be run in a separate forked process for each connection.
 
158
    #
 
159
 
 
160
    def new_client(self):
 
161
        """
 
162
        Called after a new WebSocket connection has been established.
 
163
        """
 
164
        # Checks if we receive a token, and look
 
165
        # for a valid target for it then
 
166
        if self.target_cfg:
 
167
            (self.target_host, self.target_port) = self.get_target(self.target_cfg, self.path)
 
168
 
 
169
        # Connect to the target
 
170
        if self.wrap_cmd:
 
171
            msg = "connecting to command: '%s' (port %s)" % (" ".join(self.wrap_cmd), self.target_port)
 
172
        elif self.unix_target:
 
173
            msg = "connecting to unix socket: %s" % self.unix_target
 
174
        else:
 
175
            msg = "connecting to: %s:%s" % (
 
176
                                    self.target_host, self.target_port)
 
177
 
 
178
        if self.ssl_target:
 
179
            msg += " (using SSL)"
 
180
        self.msg(msg)
 
181
 
 
182
        tsock = self.socket(self.target_host, self.target_port,
 
183
                connect=True, use_ssl=self.ssl_target, unix_socket=self.unix_target)
 
184
 
 
185
        if self.verbose and not self.daemon:
 
186
            print(self.traffic_legend)
 
187
 
 
188
        # Start proxying
 
189
        try:
 
190
            self.do_proxy(tsock)
 
191
        except:
 
192
            if tsock:
 
193
                tsock.shutdown(socket.SHUT_RDWR)
 
194
                tsock.close()
 
195
                self.vmsg("%s:%s: Closed target" %(
 
196
                    self.target_host, self.target_port))
 
197
            raise
 
198
 
 
199
    def get_target(self, target_cfg, path):
 
200
        """
 
201
        Parses the path, extracts a token, and looks for a valid
 
202
        target for that token in the configuration file(s). Sets
 
203
        target_host and target_port if successful
 
204
        """
 
205
        # The files in targets contain the lines
 
206
        # in the form of token: host:port
 
207
 
 
208
        # Extract the token parameter from url
 
209
        args = parse_qs(urlparse(path)[4]) # 4 is the query from url
 
210
 
 
211
        if not args.has_key('token') or not len(args['token']):
 
212
            raise self.EClose("Token not present")
 
213
 
 
214
        token = args['token'][0].rstrip('\n')
 
215
 
 
216
        # target_cfg can be a single config file or directory of
 
217
        # config files
 
218
        if os.path.isdir(target_cfg):
 
219
            cfg_files = [os.path.join(target_cfg, f)
 
220
                         for f in os.listdir(target_cfg)]
 
221
        else:
 
222
            cfg_files = [target_cfg]
 
223
 
 
224
        targets = {}
 
225
        for f in cfg_files:
 
226
            for line in [l.strip() for l in file(f).readlines()]:
 
227
                if line and not line.startswith('#'):
 
228
                    ttoken, target = line.split(': ')
 
229
                    targets[ttoken] = target.strip()
 
230
 
 
231
        self.vmsg("Target config: %s" % repr(targets))
 
232
 
 
233
        if targets.has_key(token):
 
234
            return targets[token].split(':')
 
235
        else:
 
236
            raise self.EClose("Token '%s' not found" % token)
 
237
 
 
238
    def do_proxy(self, target):
 
239
        """
 
240
        Proxy client WebSocket to normal target socket.
 
241
        """
 
242
        cqueue = []
 
243
        c_pend = 0
 
244
        tqueue = []
 
245
        rlist = [self.client, target]
 
246
 
 
247
        while True:
 
248
            wlist = []
 
249
 
 
250
            if tqueue: wlist.append(target)
 
251
            if cqueue or c_pend: wlist.append(self.client)
 
252
            ins, outs, excepts = select(rlist, wlist, [], 1)
 
253
            if excepts: raise Exception("Socket exception")
 
254
 
 
255
            if self.client in outs:
 
256
                # Send queued target data to the client
 
257
                c_pend = self.send_frames(cqueue)
 
258
 
 
259
                cqueue = []
 
260
 
 
261
            if self.client in ins:
 
262
                # Receive client data, decode it, and queue for target
 
263
                bufs, closed = self.recv_frames()
 
264
                tqueue.extend(bufs)
 
265
 
 
266
                if closed:
 
267
                    # TODO: What about blocking on client socket?
 
268
                    self.vmsg("%s:%s: Client closed connection" %(
 
269
                        self.target_host, self.target_port))
 
270
                    raise self.CClose(closed['code'], closed['reason'])
 
271
 
 
272
 
 
273
            if target in outs:
 
274
                # Send queued client data to the target
 
275
                dat = tqueue.pop(0)
 
276
                sent = target.send(dat)
 
277
                if sent == len(dat):
 
278
                    self.traffic(">")
 
279
                else:
 
280
                    # requeue the remaining data
 
281
                    tqueue.insert(0, dat[sent:])
 
282
                    self.traffic(".>")
 
283
 
 
284
 
 
285
            if target in ins:
 
286
                # Receive target data, encode it and queue for client
 
287
                buf = target.recv(self.buffer_size)
 
288
                if len(buf) == 0:
 
289
                    self.vmsg("%s:%s: Target closed connection" %(
 
290
                        self.target_host, self.target_port))
 
291
                    raise self.CClose(1000, "Target closed")
 
292
 
 
293
                cqueue.append(buf)
 
294
                self.traffic("{")
 
295
 
 
296
 
 
297
 
 
298
def _subprocess_setup():
 
299
    # Python installs a SIGPIPE handler by default. This is usually not what
 
300
    # non-Python successfulbprocesses expect.
 
301
    signal.signal(signal.SIGPIPE, signal.SIG_DFL)
 
302
 
 
303
 
 
304
def websockify_init():
 
305
    usage = "\n    %prog [options]"
 
306
    usage += " [source_addr:]source_port [target_addr:target_port]"
 
307
    usage += "\n    %prog [options]"
 
308
    usage += " [source_addr:]source_port -- WRAP_COMMAND_LINE"
 
309
    parser = optparse.OptionParser(usage=usage)
 
310
    parser.add_option("--verbose", "-v", action="store_true",
 
311
            help="verbose messages and per frame traffic")
 
312
    parser.add_option("--record",
 
313
            help="record sessions to FILE.[session_number]", metavar="FILE")
 
314
    parser.add_option("--daemon", "-D",
 
315
            dest="daemon", action="store_true",
 
316
            help="become a daemon (background process)")
 
317
    parser.add_option("--run-once", action="store_true",
 
318
            help="handle a single WebSocket connection and exit")
 
319
    parser.add_option("--timeout", type=int, default=0,
 
320
            help="after TIMEOUT seconds exit when not connected")
 
321
    parser.add_option("--idle-timeout", type=int, default=0,
 
322
            help="server exits after TIMEOUT seconds if there are no "
 
323
                 "active connections")
 
324
    parser.add_option("--cert", default="self.pem",
 
325
            help="SSL certificate file")
 
326
    parser.add_option("--key", default=None,
 
327
            help="SSL key file (if separate from cert)")
 
328
    parser.add_option("--ssl-only", action="store_true",
 
329
            help="disallow non-encrypted client connections")
 
330
    parser.add_option("--ssl-target", action="store_true",
 
331
            help="connect to SSL target as SSL client")
 
332
    parser.add_option("--unix-target",
 
333
            help="connect to unix socket target", metavar="FILE")
 
334
    parser.add_option("--web", default=None, metavar="DIR",
 
335
            help="run webserver on same port. Serve files from DIR.")
 
336
    parser.add_option("--wrap-mode", default="exit", metavar="MODE",
 
337
            choices=["exit", "ignore", "respawn"],
 
338
            help="action to take when the wrapped program exits "
 
339
            "or daemonizes: exit (default), ignore, respawn")
 
340
    parser.add_option("--prefer-ipv6", "-6",
 
341
            action="store_true", dest="source_is_ipv6",
 
342
            help="prefer IPv6 when resolving source_addr")
 
343
    parser.add_option("--target-config", metavar="FILE",
 
344
            dest="target_cfg",
 
345
            help="Configuration file containing valid targets "
 
346
            "in the form 'token: host:port' or, alternatively, a "
 
347
            "directory containing configuration files of this form")
 
348
    (opts, args) = parser.parse_args()
 
349
 
 
350
    # Sanity checks
 
351
    if len(args) < 2 and not (opts.target_cfg or opts.unix_target):
 
352
        parser.error("Too few arguments")
 
353
    if sys.argv.count('--'):
 
354
        opts.wrap_cmd = args[1:]
 
355
    else:
 
356
        opts.wrap_cmd = None
 
357
        if len(args) > 2:
 
358
            parser.error("Too many arguments")
 
359
 
 
360
    if not websocket.ssl and opts.ssl_target:
 
361
        parser.error("SSL target requested and Python SSL module not loaded.");
 
362
 
 
363
    if opts.ssl_only and not os.path.exists(opts.cert):
 
364
        parser.error("SSL only and %s not found" % opts.cert)
 
365
 
 
366
    # Parse host:port and convert ports to numbers
 
367
    if args[0].count(':') > 0:
 
368
        opts.listen_host, opts.listen_port = args[0].rsplit(':', 1)
 
369
        opts.listen_host = opts.listen_host.strip('[]')
 
370
    else:
 
371
        opts.listen_host, opts.listen_port = '', args[0]
 
372
 
 
373
    try:    opts.listen_port = int(opts.listen_port)
 
374
    except: parser.error("Error parsing listen port")
 
375
 
 
376
    if opts.wrap_cmd or opts.unix_target or opts.target_cfg:
 
377
        opts.target_host = None
 
378
        opts.target_port = None
 
379
    else:
 
380
        if args[1].count(':') > 0:
 
381
            opts.target_host, opts.target_port = args[1].rsplit(':', 1)
 
382
            opts.target_host = opts.target_host.strip('[]')
 
383
        else:
 
384
            parser.error("Error parsing target")
 
385
        try:    opts.target_port = int(opts.target_port)
 
386
        except: parser.error("Error parsing target port")
 
387
 
 
388
    # Create and start the WebSockets proxy
 
389
    server = WebSocketProxy(**opts.__dict__)
 
390
    server.start_server()
 
391
 
 
392
if __name__ == '__main__':
 
393
    websockify_init()