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 sys, socket, ssl, optparse
  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
  41     while True:
  42         wlist = []
  43         if tqueue: wlist.append(target)
  44         if cqueue: wlist.append(client)
  45         ins, outs, excepts = select(rlist, wlist, [], 1)
  46         if excepts: raise Exception("Socket exception")
  47
  48         if target in outs:
  49             dat = tqueue.pop(0)
  50             sent = target.send(dat)
  51             if sent == len(dat):
  52                 traffic(">")
  53             else:
  54                 tqueue.insert(0, dat[sent:])
  55                 traffic(".>")
  56             ##if rec: rec.write("Target send: %s\n" % map(ord, dat))
  57
  58         if client in outs:
  59             dat = cqueue.pop(0)
  60             sent = client.send(dat)
  61             if sent == len(dat):
  62                 traffic("<")
  63                 ##if rec: rec.write("Client send: %s ...\n" % repr(dat[0:80]))
  64                 if rec: rec.write("%s,\n" % repr(">" + dat[1:-1]))
  65             else:
  66                 cqueue.insert(0, dat[sent:])
  67                 traffic("<.")
  68                 ##if rec: rec.write("Client send partial: %s\n" % repr(dat[0:send]))
  69
  70
  71         if target in ins:
  72             buf = target.recv(buffer_size)
  73             if len(buf) == 0: raise Exception("Target closed")
  74
  75             cqueue.append(encode(buf))
  76             traffic("{")
  77             ##if rec: rec.write("Target recv (%d): %s\n" % (len(buf), map(ord, buf)))
  78
  79         if client in ins:
  80             buf = client.recv(buffer_size)
  81             if len(buf) == 0: raise Exception("Client closed")
  82
  83             if buf[-1] == '\xff':
  84                 if buf.count('\xff') > 1:
  85                     traffic(str(buf.count('\xff')))
  86                 traffic("}")
  87                 ##if rec: rec.write("Client recv (%d): %s\n" % (len(buf), repr(buf)))
  88                 if rec: rec.write("%s,\n" % repr(buf[1:-1]))
  89                 if cpartial:
  90                     tqueue.extend(decode(cpartial + buf))
  91                     cpartial = ""
  92                 else:
  93                     tqueue.extend(decode(buf))
  94             else:
  95                 traffic(".}")
  96                 ##if rec: rec.write("Client recv partial (%d): %s\n" % (len(buf), repr(buf)))
  97                 cpartial = cpartial + buf
  98
  99 def proxy_handler(client):
 100     global target_host, target_port, options, rec
 101
 102     if settings['record']:
 103         print "Opening record file: %s" % settings['record']
 104         rec = open(settings['record'], 'a')
 105
 106     print "Connecting to: %s:%s" % (target_host, target_port)
 107     tsock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 108     tsock.connect((target_host, target_port))
 109
 110     print traffic_legend
 111
 112     try:
 113         do_proxy(client, tsock)
 114     except:
 115         if tsock: tsock.close()
 116         if rec: rec.close()
 117         raise
 118
 119 if __name__ == '__main__':
 120     usage = "%prog [--record FILE]"
 121     usage += " [source_addr:]source_port target_addr:target_port"
 122     parser = optparse.OptionParser(usage=usage)
 123     parser.add_option("--record",
 124             help="record session to a file", metavar="FILE")
 125     parser.add_option("--foreground", "-f",
 126             dest="daemon", default=True, action="store_false",
 127             help="stay in foreground, do not daemonize")
 128     parser.add_option("--ssl-only", action="store_true",
 129             help="disallow non-encrypted connections")
 130     parser.add_option("--cert", default="self.pem",
 131             help="SSL certificate")
 132     (options, args) = parser.parse_args()
 133
 134     if len(args) > 2: parser.error("Too many arguments")
 135     if len(args) < 2: parser.error("Too few arguments")
 136     if args[0].count(':') > 0:
 137         host,port = args[0].split(':')
 138     else:
 139         host,port = '',args[0]
 140     if args[1].count(':') > 0:
 141         target_host,target_port = args[1].split(':')
 142     else:
 143         parser.error("Error parsing target")
 144     try:    port = int(port)
 145     except: parser.error("Error parsing listen port")
 146     try:    target_port = int(target_port)
 147     except: parser.error("Error parsing target port")
 148
 149     if options.ssl_only and not os.path.exists(options.cert):
 150         parser.error("SSL only and %s not found" % options.cert)
 151
 152     settings['listen_host'] = host
 153     settings['listen_port'] = port
 154     settings['handler'] = proxy_handler
 155     settings['cert'] = os.path.abspath(options.cert)
 156     settings['ssl_only'] = options.ssl_only
 157     settings['daemon'] = options.daemon
 158     if options.record:
 159         settings['record'] = os.path.abspath(options.record)
 160     start_server()