This commit is contained in:
Alon Bar-Lev 2013-10-15 09:18:24 -07:00
commit 50e04be7df
2 changed files with 93 additions and 55 deletions

View File

@ -16,7 +16,7 @@ as taken from http://docs.python.org/dev/library/ssl.html#certificates
'''
import os, sys, time, errno, signal, socket, traceback, select
import os, sys, time, errno, signal, socket, select, logging
import array, struct
from base64 import b64encode, b64decode
@ -58,7 +58,7 @@ for mod, msg in [('numpy', 'HyBi protocol will be slower'),
globals()[mod] = __import__(mod)
except ImportError:
globals()[mod] = None
print("WARNING: no '%s' module, %s" % (mod, msg))
self.msg("WARNING: no '%s' module, %s", mod, msg)
if multiprocessing and sys.platform == 'win32':
# make sockets pickle-able/inheritable
import multiprocessing.reduction
@ -70,6 +70,7 @@ class WebSocketServer(object):
Must be sub-classed with new_client method definition.
"""
log_prefix = "websocket"
buffer_size = 65536
server_handshake_hybi = """HTTP/1.1 101 Switching Protocols\r
@ -96,7 +97,7 @@ Sec-WebSocket-Accept: %s\r
def __init__(self, listen_host='', listen_port=None, source_is_ipv6=False,
verbose=False, cert='', key='', ssl_only=None,
daemon=False, record='', web='', file_only=False, no_parent=False,
run_once=False, timeout=0, idle_timeout=0):
run_once=False, timeout=0, idle_timeout=0, traffic=False):
# settings
self.verbose = verbose
@ -108,14 +109,18 @@ Sec-WebSocket-Accept: %s\r
self.run_once = run_once
self.timeout = timeout
self.idle_timeout = idle_timeout
self.traffic = traffic
self.launch_time = time.time()
self.ws_connection = False
self.i_am_client = False
self.handler_id = 1
self.file_only = file_only
self.no_parent = no_parent
self.logger = self.get_logger()
# Make paths settings absolute
self.cert = os.path.abspath(cert)
self.key = self.web = self.record = ''
@ -136,30 +141,36 @@ Sec-WebSocket-Accept: %s\r
raise Exception("Module 'resource' required to daemonize")
# Show configuration
print("WebSocket server settings:")
print(" - Listen on %s:%s" % (
self.listen_host, self.listen_port))
print(" - Flash security policy server")
self.msg("WebSocket server settings:")
self.msg(" - Listen on %s:%s",
self.listen_host, self.listen_port)
self.msg(" - Flash security policy server")
if self.web:
print(" - Web server. Web root: %s" % self.web)
self.msg(" - Web server. Web root: %s", self.web)
if ssl:
if os.path.exists(self.cert):
print(" - SSL/TLS support")
self.msg(" - SSL/TLS support")
if self.ssl_only:
print(" - Deny non-SSL/TLS connections")
self.msg(" - Deny non-SSL/TLS connections")
else:
print(" - No SSL/TLS support (no cert file)")
self.msg(" - No SSL/TLS support (no cert file)")
else:
print(" - No SSL/TLS support (no 'ssl' module)")
self.msg(" - No SSL/TLS support (no 'ssl' module)")
if self.daemon:
print(" - Backgrounding (daemon)")
self.msg(" - Backgrounding (daemon)")
if self.record:
print(" - Recording to '%s.*'" % self.record)
self.msg(" - Recording to '%s.*'", self.record)
#
# WebSocketServer static methods
#
@staticmethod
def get_logger():
return logging.getLogger("%s.%s" % (
WebSocketServer.log_prefix,
WebSocketServer.__class__.__name__))
@staticmethod
def socket(host, port=None, connect=False, prefer_ipv6=False, unix_socket=None, use_ssl=False):
""" Resolve a host (and optional port) to an IPv4 or IPv6
@ -251,7 +262,7 @@ Sec-WebSocket-Accept: %s\r
b = numpy.bitwise_xor(data, mask).tostring()
if plen % 4:
#print("Partial unmask")
#self.msg("Partial unmask")
mask = numpy.frombuffer(buf, dtype=numpy.dtype('B'),
offset=hlen, count=(plen % 4))
data = numpy.frombuffer(buf, dtype=numpy.dtype('B'),
@ -292,7 +303,7 @@ Sec-WebSocket-Accept: %s\r
elif payload_len >= 65536:
header = pack('>BBQ', b1, 127, payload_len)
#print("Encoded: %s" % repr(header + buf))
#self.msg("Encoded: %s", repr(header + buf))
return header + buf, len(header), 0
@ -321,6 +332,8 @@ Sec-WebSocket-Accept: %s\r
'close_code' : 1000,
'close_reason' : ''}
logger = WebSocketServer.get_logger()
blen = len(buf)
f['left'] = blen
@ -359,15 +372,16 @@ Sec-WebSocket-Accept: %s\r
f['payload'] = WebSocketServer.unmask(buf, f['hlen'],
f['length'])
else:
print("Unmasked frame: %s" % repr(buf))
self.vmsg("Unmasked frame: %s" % repr(buf))
f['payload'] = buf[(f['hlen'] + f['masked'] * 4):full_len]
if base64 and f['opcode'] in [1, 2]:
try:
f['payload'] = b64decode(f['payload'])
except:
print("Exception while b64decoding buffer: %s" %
self.warn("Exception while b64decoding buffer: %s",
repr(buf))
self.vmsg('Exception', exc_info=True)
raise
if f['opcode'] == 0x08:
@ -383,21 +397,32 @@ Sec-WebSocket-Accept: %s\r
# WebSocketServer logging/output functions
#
def traffic(self, token="."):
""" Show traffic flow in verbose mode. """
if self.verbose and not self.daemon:
def print_traffic(self, token="."):
""" Show traffic flow mode. """
if self.traffic:
sys.stdout.write(token)
sys.stdout.flush()
def msg(self, msg):
""" Output message with handler_id prefix. """
if not self.daemon:
print("% 3d: %s" % (self.handler_id, msg))
def vmsg(self, msg):
""" Same as msg() but only if verbose. """
if self.verbose:
self.msg(msg)
def log(self, lvl, msg, *args, **kwargs):
""" Wrapper around python logging """
prefix = ""
if self.i_am_client:
prefix = "% 3d: " % self.handler_id
self.logger.log(lvl, "%s%s" % (prefix, msg),
*args, **kwargs)
def msg(self, *args, **kwargs):
""" Output message with handler_id prefix. """
self.log(logging.INFO, *args, **kwargs)
def vmsg(self, *args, **kwargs):
""" Same as msg() but as debug. """
self.log(logging.DEBUG, *args, **kwargs)
def warn(self, *args, **kwargs):
""" Same as msg() but as warning. """
self.log(logging.WARN, *args, **kwargs)
#
# Main WebSocketServer methods
@ -432,9 +457,9 @@ Sec-WebSocket-Accept: %s\r
sent = self.client.send(buf)
if sent == len(buf):
self.traffic("<")
self.print_traffic("<")
else:
self.traffic("<.")
self.print_traffic("<.")
self.send_parts.insert(0, buf[sent:])
break
@ -463,11 +488,11 @@ Sec-WebSocket-Accept: %s\r
while buf:
frame = self.decode_hybi(buf, base64=self.base64)
#print("Received buf: %s, frame: %s" % (repr(buf), frame))
#self.msg("Received buf: %s, frame: %s", repr(buf), frame)
if frame['payload'] == None:
# Incomplete/partial frame
self.traffic("}.")
self.print_traffic("}.")
if frame['left'] > 0:
self.recv_part = buf[-frame['left']:]
break
@ -477,7 +502,7 @@ Sec-WebSocket-Accept: %s\r
'reason': frame['close_reason']}
break
self.traffic("}")
self.print_traffic("}")
if self.rec:
start = frame['hlen']
@ -693,6 +718,7 @@ Sec-WebSocket-Accept: %s\r
def top_new_client(self, startsock, address):
""" Do something with a WebSockets client connection. """
# Initialize per client settings
self.i_am_client = True
self.send_parts = []
self.recv_part = None
self.base64 = False
@ -718,7 +744,7 @@ Sec-WebSocket-Accept: %s\r
self.ws_connection = True
self.new_client()
except self.CClose, WebSocketServer.Terminate:
except self.CClose:
# Close the client
_, exc, _ = sys.exc_info()
if self.client:
@ -733,8 +759,7 @@ Sec-WebSocket-Accept: %s\r
except Exception:
_, exc, _ = sys.exc_info()
self.msg("handler exception: %s" % str(exc))
if self.verbose:
self.msg(traceback.format_exc())
self.vmsg("exception", exc_info=True)
finally:
if self.rec:
self.rec.write("'EOF'];\n")
@ -857,22 +882,19 @@ Sec-WebSocket-Accept: %s\r
self.handler_id += 1
except (self.Terminate, SystemExit, KeyboardInterrupt):
_, exc, _ = sys.exc_info()
print("In exit")
self.msg("In exit")
break
except Exception:
_, exc, _ = sys.exc_info()
self.msg("handler exception: %s" % str(exc))
if self.verbose:
self.msg(traceback.format_exc())
self.msg("handler exception: %s", str(exc))
self.vmsg("exception", exc_info=True)
finally:
if startsock:
startsock.close()
finally:
# Close listen port
self.vmsg("Closing socket listening at %s:%s"
% (self.listen_host, self.listen_port))
self.vmsg("Closing socket listening at %s:%s",
self.listen_host, self.listen_port)
lsock.close()
# Restore signals

View File

@ -11,7 +11,7 @@ as taken from http://docs.python.org/dev/library/ssl.html#certificates
'''
import signal, socket, optparse, time, os, sys, subprocess
import signal, socket, optparse, time, os, sys, subprocess, logging
from select import select
import websocket
try:
@ -86,7 +86,7 @@ Traffic Legend:
websocket.WebSocketServer.__init__(self, *args, **kwargs)
def run_wrap_cmd(self):
print("Starting '%s'" % " ".join(self.wrap_cmd))
self.msg("Starting '%s'", " ".join(self.wrap_cmd))
self.wrap_times.append(time.time())
self.wrap_times.pop(0)
self.cmd = subprocess.Popen(
@ -116,7 +116,7 @@ Traffic Legend:
if self.ssl_target:
msg += " (using SSL)"
print(msg + "\n")
self.msg("%s", msg)
if self.wrap_cmd:
self.run_wrap_cmd()
@ -142,7 +142,7 @@ Traffic Legend:
if (now - avg) < 10:
# 3 times in the last 10 seconds
if self.spawn_message:
print("Command respawning too fast")
self.warn("Command respawning too fast")
self.spawn_message = False
else:
self.run_wrap_cmd()
@ -182,8 +182,7 @@ Traffic Legend:
tsock = self.socket(self.target_host, self.target_port,
connect=True, use_ssl=self.ssl_target, unix_socket=self.unix_target)
if self.verbose and not self.daemon:
print(self.traffic_legend)
self.print_traffic(self.traffic_legend)
# Start proxying
try:
@ -275,11 +274,11 @@ Traffic Legend:
dat = tqueue.pop(0)
sent = target.send(dat)
if sent == len(dat):
self.traffic(">")
self.print_traffic(">")
else:
# requeue the remaining data
tqueue.insert(0, dat[sent:])
self.traffic(".>")
self.print_traffic(".>")
if target in ins:
@ -291,7 +290,7 @@ Traffic Legend:
raise self.CClose(1000, "Target closed")
cqueue.append(buf)
self.traffic("{")
self.print_traffic("{")
@ -301,14 +300,28 @@ def _subprocess_setup():
signal.signal(signal.SIGPIPE, signal.SIG_DFL)
def logger_init():
logger = logging.getLogger(WebSocketProxy.log_prefix)
logger.propagate = False
logger.setLevel(logging.INFO)
h = logging.StreamHandler()
h.setLevel(logging.DEBUG)
h.setFormatter(logging.Formatter("%(message)s"))
logger.addHandler(h)
def websockify_init():
logger_init()
usage = "\n %prog [options]"
usage += " [source_addr:]source_port [target_addr:target_port]"
usage += "\n %prog [options]"
usage += " [source_addr:]source_port -- WRAP_COMMAND_LINE"
parser = optparse.OptionParser(usage=usage)
parser.add_option("--verbose", "-v", action="store_true",
help="verbose messages and per frame traffic")
help="verbose messages")
parser.add_option("--traffic", action="store_true",
help="per frame traffic")
parser.add_option("--record",
help="record sessions to FILE.[session_number]", metavar="FILE")
parser.add_option("--daemon", "-D",
@ -347,6 +360,9 @@ def websockify_init():
"directory containing configuration files of this form")
(opts, args) = parser.parse_args()
if opts.verbose:
logging.getLogger(WebSocketProxy.log_prefix).setLevel(logging.DEBUG)
# Sanity checks
if len(args) < 2 and not (opts.target_cfg or opts.unix_target):
parser.error("Too few arguments")