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 Exception("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 Exception("Client closed")
  84
  85             if buf == '\xff\x00':
  86                 raise Exception("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
 105
 106     if settings['record']:
 107         print "Opening record file: %s" % settings['record']
 108         rec = open(settings['record'], 'w+')
 109         rec.write("var VNC_frame_data = [\n")
 110
 111     print "Connecting to: %s:%s" % (target_host, target_port)
 112     tsock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 113     tsock.connect((target_host, target_port))
 114
 115     print traffic_legend
 116
 117     try:
 118         do_proxy(client, tsock)
 119     except:
 120         if tsock: tsock.close()
 121         if rec:
 122             rec.write("'EOF']\n")
 123             rec.close()
 124         raise
 125
 126 if __name__ == '__main__':
 127     usage = "%prog [--record FILE]"
 128     usage += " [source_addr:]source_port target_addr:target_port"
 129     parser = optparse.OptionParser(usage=usage)
 130     parser.add_option("--record",
 131             help="record session to a file", metavar="FILE")
 132     parser.add_option("--foreground", "-f",
 133             dest="daemon", default=True, action="store_false",
 134             help="stay in foreground, do not daemonize")
 135     parser.add_option("--ssl-only", action="store_true",
 136             help="disallow non-encrypted connections")
 137     parser.add_option("--cert", default="self.pem",
 138             help="SSL certificate")
 139     (options, args) = parser.parse_args()
 140
 141     if len(args) > 2: parser.error("Too many arguments")
 142     if len(args) < 2: parser.error("Too few arguments")
 143     if args[0].count(':') > 0:
 144         host,port = args[0].split(':')
 145     else:
 146         host,port = '',args[0]
 147     if args[1].count(':') > 0:
 148         target_host,target_port = args[1].split(':')
 149     else:
 150         parser.error("Error parsing target")
 151     try:    port = int(port)
 152     except: parser.error("Error parsing listen port")
 153     try:    target_port = int(target_port)
 154     except: parser.error("Error parsing target port")
 155
 156     if options.ssl_only and not os.path.exists(options.cert):
 157         parser.error("SSL only and %s not found" % options.cert)
 158
 159     settings['listen_host'] = host
 160     settings['listen_port'] = port
 161     settings['handler'] = proxy_handler
 162     settings['cert'] = os.path.abspath(options.cert)
 163     settings['ssl_only'] = options.ssl_only
 164     settings['daemon'] = options.daemon
 165     if options.record:
 166         settings['record'] = os.path.abspath(options.record)
 167     start_server()