utils/wsproxy.py

   1 #!/usr/bin/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
  15 from select import select
  16 from websocket import *
  17
  18 buffer_size = 65536
  19 rec = None
  20
  21 traffic_legend = """
  22 Traffic Legend:
  23     }  - Client receive
  24     }. - Client receive partial
  25     {  - Target receive
  26
  27     >  - Target send
  28     >. - Target send partial
  29     <  - Client send
  30     <. - Client send partial
  31 """
  32
  33 def do_proxy(client, target):
  34     """ Proxy WebSocket to normal socket. """
  35     global rec
  36     cqueue = []
  37     cpartial = ""
  38     tqueue = []
  39     rlist = [client, target]
  40     tstart = int(time.time()*1000)
  41
  42     while True:
  43         wlist = []
  44         tdelta = int(time.time()*1000) - tstart
  45         if tqueue: wlist.append(target)
  46         if cqueue: wlist.append(client)
  47         ins, outs, excepts = select(rlist, wlist, [], 1)
  48         if excepts: raise Exception("Socket exception")
  49
  50         if target in outs:
  51             dat = tqueue.pop(0)
  52             sent = target.send(dat)
  53             if sent == len(dat):
  54                 traffic(">")
  55             else:
  56                 tqueue.insert(0, dat[sent:])
  57                 traffic(".>")
  58             ##if rec: rec.write("Target send: %s\n" % map(ord, dat))
  59
  60         if client in outs:
  61             dat = cqueue.pop(0)
  62             sent = client.send(dat)
  63             if sent == len(dat):
  64                 traffic("<")
  65                 ##if rec: rec.write("Client send: %s ...\n" % repr(dat[0:80]))
  66                 if rec: rec.write("%s,\n" % repr("{%s{" % tdelta + dat[1:-1]))
  67             else:
  68                 cqueue.insert(0, dat[sent:])
  69                 traffic("<.")
  70                 ##if rec: rec.write("Client send partial: %s\n" % repr(dat[0:send]))
  71
  72
  73         if target in ins:
  74             buf = target.recv(buffer_size)
  75             if len(buf) == 0: raise EClose("Target closed")
  76
  77             cqueue.append(encode(buf))
  78             traffic("{")
  79             ##if rec: rec.write("Target recv (%d): %s\n" % (len(buf), map(ord, buf)))
  80
  81         if client in ins:
  82             buf = client.recv(buffer_size)
  83             if len(buf) == 0: raise EClose("Client closed")
  84
  85             if buf == '\xff\x00':
  86                 raise EClose("Client sent orderly close frame")
  87             elif buf[-1] == '\xff':
  88                 if buf.count('\xff') > 1:
  89                     traffic(str(buf.count('\xff')))
  90                 traffic("}")
  91                 ##if rec: rec.write("Client recv (%d): %s\n" % (len(buf), repr(buf)))
  92                 if rec: rec.write("%s,\n" % (repr("}%s}" % tdelta + buf[1:-1])))
  93                 if cpartial:
  94                     tqueue.extend(decode(cpartial + buf))
  95                     cpartial = ""
  96                 else:
  97                     tqueue.extend(decode(buf))
  98             else:
  99                 traffic(".}")
 100                 ##if rec: rec.write("Client recv partial (%d): %s\n" % (len(buf), repr(buf)))
 101                 cpartial = cpartial + buf
 102
 103 def proxy_handler(client):
 104     global target_host, target_port, options, rec, fname
 105
 106     if settings['record']:
 107         fname = "%s.%s" % (settings['record'],
 108                             settings['handler_id'])
 109         handler_msg("opening record file: %s" % fname)
 110         rec = open(fname, 'w+')
 111         rec.write("var VNC_frame_data = [\n")
 112
 113     handler_msg("connecting to: %s:%s" % (target_host, target_port))
 114     tsock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 115     tsock.connect((target_host, target_port))
 116
 117     if settings['verbose'] and not settings['daemon']:
 118         print traffic_legend
 119
 120     try:
 121         do_proxy(client, tsock)
 122     except:
 123         if tsock: tsock.close()
 124         if rec:
 125             rec.write("'EOF']\n")
 126             rec.close()
 127         raise
 128
 129 if __name__ == '__main__':
 130     usage = "%prog [--record FILE]"
 131     usage += " [source_addr:]source_port target_addr:target_port"
 132     parser = optparse.OptionParser(usage=usage)
 133     parser.add_option("--verbose", "-v", action="store_true",
 134             help="verbose messages and per frame traffic")
 135     parser.add_option("--record",
 136             help="record sessions to FILE.[session_number]", metavar="FILE")
 137     parser.add_option("--foreground", "-f",
 138             dest="daemon", default=True, action="store_false",
 139             help="stay in foreground, do not daemonize")
 140     parser.add_option("--cert", default="self.pem",
 141             help="SSL certificate file")
 142     parser.add_option("--key", default=None,
 143             help="SSL key file (if separate from cert)")
 144     parser.add_option("--ssl-only", action="store_true",
 145             help="disallow non-encrypted connections")
 146     (options, args) = parser.parse_args()
 147
 148     if len(args) > 2: parser.error("Too many arguments")
 149     if len(args) < 2: parser.error("Too few arguments")
 150     if args[0].count(':') > 0:
 151         host,port = args[0].split(':')
 152     else:
 153         host,port = '',args[0]
 154     if args[1].count(':') > 0:
 155         target_host,target_port = args[1].split(':')
 156     else:
 157         parser.error("Error parsing target")
 158     try:    port = int(port)
 159     except: parser.error("Error parsing listen port")
 160     try:    target_port = int(target_port)
 161     except: parser.error("Error parsing target port")
 162
 163     if options.ssl_only and not os.path.exists(options.cert):
 164         parser.error("SSL only and %s not found" % options.cert)
 165     elif not os.path.exists(options.cert):
 166         print "Warning: %s not found" % options.cert
 167
 168     settings['verbose'] = options.verbose
 169     settings['listen_host'] = host
 170     settings['listen_port'] = port
 171     settings['handler'] = proxy_handler
 172     settings['cert'] = os.path.abspath(options.cert)
 173     if options.key:
 174         settings['key'] = os.path.abspath(options.key)
 175     settings['ssl_only'] = options.ssl_only
 176     settings['daemon'] = options.daemon
 177     if options.record:
 178         settings['record'] = os.path.abspath(options.record)
 179     start_server()