crabapples
/
websockify
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Merge pull request #100 from alonbl/log 

websocket: use python logging module
This commit is contained in:
Joel Martin 2013-10-15 09:41:35 -07:00
parent 4459824cc8 8a0a47223d
commit 477dce6cf8
2 changed files with 93 additions and 55 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

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