Statistics
| Branch: | Tag: | Revision:

one / src / sunstone / public / vendor / noVNC / utils / websockify @ 5fcb9de8

History | View | Annotate | Download (11 KB)

1
#!/usr/bin/env python
2

    
3
'''
4
A WebSocket to TCP socket proxy with support for "wss://" encryption.
5
Copyright 2010 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 socket, optparse, time, os, sys, subprocess
15
from select import select
16
from websocket import WebSocketServer
17

    
18
class WebSocketProxy(WebSocketServer):
19
    """
20
    Proxy traffic to and from a WebSockets client to a normal TCP
21
    socket server target. All traffic to/from the client is base64
22
    encoded/decoded to allow binary data to be sent/received to/from
23
    the target.
24
    """
25

    
26
    buffer_size = 65536
27

    
28
    traffic_legend = """
29
Traffic Legend:
30
    }  - Client receive
31
    }. - Client receive partial
32
    {  - Target receive
33

    
34
    >  - Target send
35
    >. - Target send partial
36
    <  - Client send
37
    <. - Client send partial
38
"""
39

    
40
    def __init__(self, *args, **kwargs):
41
        # Save off proxy specific options
42
        self.target_host   = kwargs.pop('target_host')
43
        self.target_port   = kwargs.pop('target_port')
44
        self.wrap_cmd      = kwargs.pop('wrap_cmd')
45
        self.wrap_mode     = kwargs.pop('wrap_mode')
46
        # Last 3 timestamps command was run
47
        self.wrap_times    = [0, 0, 0]
48

    
49
        if self.wrap_cmd:
50
            rebinder_path = ['./', os.path.dirname(sys.argv[0])]
51
            self.rebinder = None
52

    
53
            for rdir in rebinder_path:
54
                rpath = os.path.join(rdir, "rebind.so")
55
                if os.path.exists(rpath):
56
                    self.rebinder = rpath
57
                    break
58

    
59
            if not self.rebinder:
60
                raise Exception("rebind.so not found, perhaps you need to run make")
61

    
62
            self.target_host = "127.0.0.1"  # Loopback
63
            # Find a free high port
64
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
65
            sock.bind(('', 0))
66
            self.target_port = sock.getsockname()[1]
67
            sock.close()
68

    
69
            os.environ.update({
70
                "LD_PRELOAD": self.rebinder,
71
                "REBIND_OLD_PORT": str(kwargs['listen_port']),
72
                "REBIND_NEW_PORT": str(self.target_port)})
73

    
74
        WebSocketServer.__init__(self, *args, **kwargs)
75

    
76
    def run_wrap_cmd(self):
77
        print "Starting '%s'" % " ".join(self.wrap_cmd)
78
        self.wrap_times.append(time.time())
79
        self.wrap_times.pop(0)
80
        self.cmd = subprocess.Popen(
81
                self.wrap_cmd, env=os.environ)
82
        self.spawn_message = True
83

    
84
    def started(self):
85
        """
86
        Called after Websockets server startup (i.e. after daemonize)
87
        """
88
        # Need to call wrapped command after daemonization so we can
89
        # know when the wrapped command exits
90
        if self.wrap_cmd:
91
            print "  - proxying from %s:%s to '%s' (port %s)\n" % (
92
                    self.listen_host, self.listen_port,
93
                    " ".join(self.wrap_cmd), self.target_port)
94
            self.run_wrap_cmd()
95
        else:
96
            print "  - proxying from %s:%s to %s:%s\n" % (
97
                    self.listen_host, self.listen_port,
98
                    self.target_host, self.target_port)
99

    
100
    def poll(self):
101
        # If we are wrapping a command, check it's status
102

    
103
        if self.wrap_cmd and self.cmd:
104
            ret = self.cmd.poll()
105
            if ret != None:
106
                self.vmsg("Wrapped command exited (or daemon). Returned %s" % ret)
107
                self.cmd = None
108

    
109
        if self.wrap_cmd and self.cmd == None:
110
            # Response to wrapped command being gone
111
            if self.wrap_mode == "ignore":
112
                pass
113
            elif self.wrap_mode == "exit":
114
                sys.exit(ret)
115
            elif self.wrap_mode == "respawn":
116
                now = time.time()
117
                avg = sum(self.wrap_times)/len(self.wrap_times)
118
                if (now - avg) < 10:
119
                    # 3 times in the last 10 seconds
120
                    if self.spawn_message:
121
                        print "Command respawning too fast"
122
                        self.spawn_message = False
123
                else:
124
                    self.run_wrap_cmd()
125

    
126
    # 
127
    # Routines above this point are run in the master listener
128
    # process.
129
    #
130

    
131
    #
132
    # Routines below this point are connection handler routines and
133
    # will be run in a separate forked process for each connection.
134
    #
135

    
136
    def new_client(self, client):
137
        """
138
        Called after a new WebSocket connection has been established.
139
        """
140

    
141
        self.rec = None
142
        if self.record:
143
            # Record raw frame data as a JavaScript compatible file
144
            fname = "%s.%s" % (self.record,
145
                                self.handler_id)
146
            self.msg("opening record file: %s" % fname)
147
            self.rec = open(fname, 'w+')
148
            self.rec.write("var VNC_frame_data = [\n")
149

    
150
        # Connect to the target
151
        self.msg("connecting to: %s:%s" % (
152
                 self.target_host, self.target_port))
153
        tsock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
154
        tsock.connect((self.target_host, self.target_port))
155

    
156
        if self.verbose and not self.daemon:
157
            print self.traffic_legend
158

    
159
        # Stat proxying
160
        try:
161
            self.do_proxy(client, tsock)
162
        except:
163
            if tsock:
164
                tsock.close()
165
                self.vmsg("%s:%s: Target closed" %(
166
                    self.target_host, self.target_port))
167
            if self.rec:
168
                self.rec.write("'EOF']\n")
169
                self.rec.close()
170
            raise
171

    
172
    def do_proxy(self, client, target):
173
        """
174
        Proxy client WebSocket to normal target socket.
175
        """
176
        cqueue = []
177
        cpartial = ""
178
        tqueue = []
179
        rlist = [client, target]
180
        tstart = int(time.time()*1000)
181

    
182
        while True:
183
            wlist = []
184
            tdelta = int(time.time()*1000) - tstart
185

    
186
            if tqueue: wlist.append(target)
187
            if cqueue: wlist.append(client)
188
            ins, outs, excepts = select(rlist, wlist, [], 1)
189
            if excepts: raise Exception("Socket exception")
190

    
191
            if target in outs:
192
                # Send queued client data to the target
193
                dat = tqueue.pop(0)
194
                sent = target.send(dat)
195
                if sent == len(dat):
196
                    self.traffic(">")
197
                else:
198
                    # requeue the remaining data
199
                    tqueue.insert(0, dat[sent:])
200
                    self.traffic(".>")
201

    
202
            if client in outs:
203
                # Send queued target data to the client
204
                dat = cqueue.pop(0)
205
                sent = client.send(dat)
206
                if sent == len(dat):
207
                    self.traffic("<")
208
                    if self.rec:
209
                        self.rec.write("%s,\n" %
210
                                repr("{%s{" % tdelta + dat[1:-1]))
211
                else:
212
                    cqueue.insert(0, dat[sent:])
213
                    self.traffic("<.")
214

    
215

    
216
            if target in ins:
217
                # Receive target data, encode it and queue for client
218
                buf = target.recv(self.buffer_size)
219
                if len(buf) == 0: raise self.EClose("Target closed")
220

    
221
                cqueue.append(self.encode(buf))
222
                self.traffic("{")
223

    
224
            if client in ins:
225
                # Receive client data, decode it, and queue for target
226
                buf = client.recv(self.buffer_size)
227
                if len(buf) == 0: raise self.EClose("Client closed")
228

    
229
                if buf == '\xff\x00':
230
                    raise self.EClose("Client sent orderly close frame")
231
                elif buf[-1] == '\xff':
232
                    if buf.count('\xff') > 1:
233
                        self.traffic(str(buf.count('\xff')))
234
                    self.traffic("}")
235
                    if self.rec:
236
                        self.rec.write("%s,\n" %
237
                                (repr("}%s}" % tdelta + buf[1:-1])))
238
                    if cpartial:
239
                        # Prepend saved partial and decode frame(s)
240
                        tqueue.extend(self.decode(cpartial + buf))
241
                        cpartial = ""
242
                    else:
243
                        # decode frame(s)
244
                        tqueue.extend(self.decode(buf))
245
                else:
246
                    # Save off partial WebSockets frame
247
                    self.traffic(".}")
248
                    cpartial = cpartial + buf
249

    
250
if __name__ == '__main__':
251
    usage = "\n    %prog [options]"
252
    usage += " [source_addr:]source_port target_addr:target_port"
253
    usage += "\n    %prog [options]"
254
    usage += " [source_addr:]source_port -- WRAP_COMMAND_LINE"
255
    parser = optparse.OptionParser(usage=usage)
256
    parser.add_option("--verbose", "-v", action="store_true",
257
            help="verbose messages and per frame traffic")
258
    parser.add_option("--record",
259
            help="record sessions to FILE.[session_number]", metavar="FILE")
260
    parser.add_option("--daemon", "-D",
261
            dest="daemon", action="store_true",
262
            help="become a daemon (background process)")
263
    parser.add_option("--cert", default="self.pem",
264
            help="SSL certificate file")
265
    parser.add_option("--key", default=None,
266
            help="SSL key file (if separate from cert)")
267
    parser.add_option("--ssl-only", action="store_true",
268
            help="disallow non-encrypted connections")
269
    parser.add_option("--web", default=None, metavar="DIR",
270
            help="run webserver on same port. Serve files from DIR.")
271
    parser.add_option("--wrap-mode", default="exit", metavar="MODE",
272
            choices=["exit", "ignore", "respawn"],
273
            help="action to take when the wrapped program exits "
274
            "or daemonizes: exit (default), ignore, respawn")
275
    (opts, args) = parser.parse_args()
276

    
277
    # Sanity checks
278
    if len(args) < 2:
279
        parser.error("Too few arguments")
280
    if sys.argv.count('--'):
281
        opts.wrap_cmd = args[1:]
282
    else:
283
        opts.wrap_cmd = None
284
        if len(args) > 2:
285
            parser.error("Too many arguments")
286

    
287
    if opts.ssl_only and not os.path.exists(opts.cert):
288
        parser.error("SSL only and %s not found" % opts.cert)
289

    
290
    # Parse host:port and convert ports to numbers
291
    if args[0].count(':') > 0:
292
        opts.listen_host, opts.listen_port = args[0].split(':')
293
    else:
294
        opts.listen_host, opts.listen_port = '', args[0]
295
    try:    opts.listen_port = int(opts.listen_port)
296
    except: parser.error("Error parsing listen port")
297

    
298
    if opts.wrap_cmd:
299
        opts.target_host = None
300
        opts.target_port = None
301
    else:
302
        if args[1].count(':') > 0:
303
            opts.target_host, opts.target_port = args[1].split(':')
304
        else:
305
            parser.error("Error parsing target")
306
        try:    opts.target_port = int(opts.target_port)
307
        except: parser.error("Error parsing target port")
308

    
309
    # Create and start the WebSockets proxy
310
    server = WebSocketProxy(**opts.__dict__)
311
    server.start_server()