From d023c45817988a4b2114b5197c7bf3310edec0ea Mon Sep 17 00:00:00 2001 From: Tomas Edwardsson Date: Fri, 3 Aug 2012 09:53:41 +0000 Subject: [PATCH] Moved log init and removed msg and vmsg msg and vmsg now uses self.log.* Log initialization now in websockify --- websocket.py | 73 ++++++++++++++-------------------------------------- websockify | 42 +++++++++++++++++++++++++----- 2 files changed, 55 insertions(+), 60 deletions(-) diff --git a/websocket.py b/websocket.py index 20709fe..9145cc9 100644 --- a/websocket.py +++ b/websocket.py @@ -17,7 +17,6 @@ as taken from http://docs.python.org/dev/library/ssl.html#certificates ''' import os, sys, time, errno, signal, socket, traceback, select -import warnings, logging import array, struct from cgi import parse_qsl from base64 import b64encode, b64decode @@ -127,28 +126,6 @@ Sec-WebSocket-Accept: %s\r if self.web: os.chdir(self.web) - # Handle logging - if loglevel.isdigit == False: - raise Exception("Invalid loglevel specified, must be 0-5") - - log_levels = { - '0': None, - '1': logging.CRITICAL, - '2': logging.ERROR, - '3': logging.WARNING, - '4': logging.INFO, - '5': logging.DEBUG, - } - if int(loglevel) > 5: loglevel = 5 - - logformat = '%(asctime)s %(levelname)s, %(message)s' - if logfile == None: - logging.basicConfig(format=logformat) - else: - logging.basicConfig(format=logformat, filename=logfile) - self.log = logging.getLogger('websocket') - self.log.setLevel(log_levels[loglevel]) - # Sanity checks if not ssl and self.ssl_only: raise Exception("No 'ssl' module and SSL-only specified") @@ -423,16 +400,6 @@ Sec-WebSocket-Accept: %s\r sys.stdout.write(token) sys.stdout.flush() - def msg(self, msg): - """ Output message with handler_id prefix. """ - if not self.daemon: - self.log.debug("% 3d: %s" % (self.handler_id, msg)) - - def vmsg(self, msg): - """ Same as msg() but only if verbose. """ - if self.verbose: - self.msg(msg) - # # Main WebSocketServer methods # @@ -713,14 +680,14 @@ Sec-WebSocket-Accept: %s\r if 'base64' in protocols: response += "%sWebSocket-Protocol: base64\r\n" % pre else: - self.msg("Warning: client does not report 'base64' protocol support") + self.log.warning("Client does not report 'base64' protocol support") response += "\r\n" + trailer - self.msg("%s: %s WebSocket connection" % (address[0], stype)) - self.msg("%s: Version %s, base64: '%s'" % (address[0], + self.log.info("%s: %s WebSocket connection" % (address[0], stype)) + self.log.debug("%s: Version %s, base64: '%s'" % (address[0], self.version, self.base64)) if self.path != '/': - self.msg("%s: Path: '%s'" % (address[0], self.path)) + self.log.debug("%s: Path: '%s'" % (address[0], self.path)) # Send server WebSockets handshake response @@ -736,7 +703,7 @@ Sec-WebSocket-Accept: %s\r # def started(self): """ Called after WebSockets startup """ - self.vmsg("WebSockets server started") + self.log.info("WebSockets server started") def poll(self): """ Run periodically while waiting for connections. """ @@ -745,17 +712,17 @@ Sec-WebSocket-Accept: %s\r def fallback_SIGCHLD(self, sig, stack): # Reap zombies when using os.fork() (python 2.4) - self.vmsg("Got SIGCHLD, reaping zombies") + self.log.debug("Got SIGCHLD, reaping zombies") try: result = os.waitpid(-1, os.WNOHANG) while result[0]: - self.vmsg("Reaped child process %s" % result[0]) + self.log.debug("Reaped child process %s" % result[0]) result = os.waitpid(-1, os.WNOHANG) except (OSError): pass def do_SIGINT(self, sig, stack): - self.msg("Got SIGINT, exiting") + self.log.fatal("Got SIGINT, exiting") sys.exit(0) def top_new_client(self, startsock, address): @@ -776,7 +743,7 @@ Sec-WebSocket-Accept: %s\r # Record raw frame data as JavaScript array fname = "%s.%s" % (self.record, self.handler_id) - self.msg("opening record file: %s" % fname) + self.log.debug("opening record file: %s" % fname) self.rec = open(fname, 'w+') self.rec.write("var VNC_frame_data = [\n") @@ -786,12 +753,11 @@ Sec-WebSocket-Accept: %s\r _, exc, _ = sys.exc_info() # Connection was not a WebSockets connection if exc.args[0]: - self.msg("%s: %s" % (address[0], exc.args[0])) + self.log.debug("%s: %s" % (address[0], exc.args[0])) except Exception: _, exc, _ = sys.exc_info() - self.msg("handler exception: %s" % str(exc)) - if self.verbose: - self.msg(traceback.format_exc()) + self.log.error("handler exception: %s" % str(exc)) + self.log.debug(traceback.format_exc()) finally: if self.rec: self.rec.write("'EOF']\n") @@ -833,7 +799,7 @@ Sec-WebSocket-Accept: %s\r time_elapsed = time.time() - self.launch_time if self.timeout and time_elapsed > self.timeout: - self.msg('listener exit due to --timeout %s' + self.log.info('listener exit due to --timeout %s' % self.timeout) break @@ -855,7 +821,7 @@ Sec-WebSocket-Accept: %s\r else: err = exc[0] if err == errno.EINTR: - self.vmsg("Ignoring interrupted syscall") + self.log.debug("Ignoring interrupted syscall") continue else: raise @@ -864,11 +830,11 @@ Sec-WebSocket-Accept: %s\r # Run in same process if run_once self.top_new_client(startsock, address) if self.ws_connection : - self.msg('%s: exiting due to --run-once' + self.log.info('%s: exiting due to --run-once' % address[0]) break elif multiprocessing: - self.vmsg('%s: new handler Process' % address[0]) + self.log.debug('%s: new handler Process' % address[0]) p = multiprocessing.Process( target=self.top_new_client, args=(startsock, address)) @@ -876,7 +842,7 @@ Sec-WebSocket-Accept: %s\r # child will not return else: # python 2.4 - self.vmsg('%s: forking handler' % address[0]) + self.log.debug('%s: forking handler' % address[0]) pid = os.fork() if pid == 0: # child handler process @@ -896,9 +862,8 @@ Sec-WebSocket-Accept: %s\r break except Exception: _, exc, _ = sys.exc_info() - self.msg("handler exception: %s" % str(exc)) - if self.verbose: - self.msg(traceback.format_exc()) + self.log.info("handler exception: %s" % str(exc)) + self.log.debug(traceback.format_exc()) finally: if startsock: diff --git a/websockify b/websockify index 6923d4b..a40ba55 100755 --- a/websockify +++ b/websockify @@ -14,6 +14,7 @@ as taken from http://docs.python.org/dev/library/ssl.html#certificates import socket, optparse, time, os, sys, subprocess from select import select from websocket import WebSocketServer +import warnings, logging class WebSocketProxy(WebSocketServer): """ @@ -37,6 +38,7 @@ Traffic Legend: <. - Client send partial """ + #noinspection PyUnresolvedReferences def __init__(self, *args, **kwargs): # Save off proxy specific options self.target_host = kwargs.pop('target_host') @@ -45,7 +47,11 @@ Traffic Legend: self.wrap_mode = kwargs.pop('wrap_mode') # Last 3 timestamps command was run self.wrap_times = [0, 0, 0] + self.loglevel = kwargs.pop('loglevel') + self.logfile = kwargs.pop('logfile') + + self.open_log() if self.wrap_cmd: rebinder_path = ['./', os.path.dirname(sys.argv[0])] self.rebinder = None @@ -102,11 +108,11 @@ Traffic Legend: if self.wrap_cmd and self.cmd: ret = self.cmd.poll() - if ret != None: - self.vmsg("Wrapped command exited (or daemon). Returned %s" % ret) + if ret is not None: + self.log.debug("Wrapped command exited (or daemon). Returned %s" % ret) self.cmd = None - if self.wrap_cmd and self.cmd == None: + if self.wrap_cmd and self.cmd is None: # Response to wrapped command being gone if self.wrap_mode == "ignore": pass @@ -139,7 +145,7 @@ Traffic Legend: """ # Connect to the target - self.msg("connecting to: %s:%s" % ( + self.log.debug("connecting to: %s:%s" % ( self.target_host, self.target_port)) tsock = self.socket(self.target_host, self.target_port, connect=True) @@ -154,10 +160,34 @@ Traffic Legend: if tsock: tsock.shutdown(socket.SHUT_RDWR) tsock.close() - self.vmsg("%s:%s: Target closed" %( + self.log.debug("%s:%s: Target closed" %( self.target_host, self.target_port)) raise + def open_log(self): + # Handle logging + if not self.loglevel.isdigit: + raise Exception("Invalid loglevel specified, must be 0-5") + + loglevels = { + '0': None, + '1': logging.CRITICAL, + '2': logging.ERROR, + '3': logging.WARNING, + '4': logging.INFO, + '5': logging.DEBUG, + } + if int(self.loglevel) > 5: self.loglevel = 5 + + logformat = '%(asctime)s %(levelname)s, %(message)s' + if self.logfile is None: + logging.basicConfig(format=logformat) + else: + logging.basicConfig(format=logformat, filename=self.logfile) + self.log = logging.getLogger('websocket') + self.log.setLevel(loglevels[self.loglevel]) + + def do_proxy(self, target): """ Proxy client WebSocket to normal target socket. @@ -190,7 +220,7 @@ Traffic Legend: if target in ins: # Receive target data, encode it and queue for client buf = target.recv(self.buffer_size) - if len(buf) == 0: raise self.EClose("Target closed") + if not len(buf): raise self.EClose("Target closed") cqueue.append(buf) self.traffic("{")