From cf53c18d375cf40a71294771b2eec2ec0a6ca54e Mon Sep 17 00:00:00 2001
From: liris <liris.pp@gmail.com>
Date: Wed, 29 Feb 2012 10:36:33 +0900
Subject: [PATCH] - wsdump support verbose mode.   - show network raw data.   -
 show opcodes

---
 bin/wsdump.py | 29 +++++++++++++++++++++++++----
 websocket.py  | 23 +++++++++++++++++++----
 2 files changed, 44 insertions(+), 8 deletions(-)

diff --git a/bin/wsdump.py b/bin/wsdump.py
index df6e4ab..02b40af 100644
--- a/bin/wsdump.py
+++ b/bin/wsdump.py
@@ -14,11 +14,24 @@ except:
 OPCODE_DATA = (websocket.ABNF.OPCODE_TEXT, websocket.ABNF.OPCODE_BINARY)
 ENCODING = getattr(sys.stdin, "encoding", "").lower()
 
+class VAction(argparse.Action):
+    def __call__(self, parser, args, values, option_string=None):
+        if values==None:
+            values = "1"
+        try:
+            values = int(values)
+        except ValueError:
+            values = values.count("v")+1
+        setattr(args, self.dest, values)
 
 def parse_args():
     parser = argparse.ArgumentParser(description="WebSocket Simple Dump Tool")
     parser.add_argument("url", metavar="ws_url",
                         help="websocket url. ex. ws://echo.websocket.org/")
+    parser.add_argument("-v", "--verbose", default=0, nargs='?', action=VAction,
+                        dest="verbose",
+                        help="set verbose mode. If set to 1, show opcode. "
+                        "If set to 2, enable to trace  websocket module")
 
     return parser.parse_args()
 
@@ -45,9 +58,11 @@ class InteractiveConsole(code.InteractiveConsole):
 def main():
     args = parse_args()
     console = InteractiveConsole()
-    print "Press Ctrl+C to quit"
     ws = websocket.create_connection(args.url)
-
+    if args.verbose > 1:
+        websocket.enableTrace(True)
+    print "Press Ctrl+C to quit"
+    
     def recv():
         frame = ws.recv_frame()
         if not frame:
@@ -66,8 +81,14 @@ def main():
     def recv_ws():
         while True:
             opcode, data = recv()
-            if opcode in OPCODE_DATA:
-                console.write("< %s" % data)
+            msg = None
+            if not args.verbose and opcode in OPCODE_DATA:
+                msg = "< %s" % data
+            elif args.verbose:
+                msg = "< %s: %s" % (websocket.ABNF.OPCODE_MAP.get(opcode), data)
+
+            if msg:
+                console.write(msg)
 
     thread = threading.Thread(target=recv_ws)
     thread.daemon = True
diff --git a/websocket.py b/websocket.py
index 6561f30..b0d7b81 100644
--- a/websocket.py
+++ b/websocket.py
@@ -210,6 +210,15 @@ class ABNF(object):
     OPCODES = (OPCODE_TEXT, OPCODE_BINARY, OPCODE_CLOSE,
                 OPCODE_PING, OPCODE_PONG)
 
+    # opcode human readable string
+    OPCODE_MAP = {
+        OPCODE_TEXT: "text",
+        OPCODE_BINARY: "binary",
+        OPCODE_CLOSE: "close",
+        OPCODE_PING: "ping",
+        OPCODE_PONG: "pong"
+        }
+
     # data length threashold.
     LENGTH_7  = 0x7d
     LENGTH_16 = 1 << 16
@@ -550,16 +559,22 @@ class WebSocket(object):
         mask = b2 >> 7 & 1
         length = b2 & 0x7f
 
+        length_data = ""
         if length == 0x7e:
-            l = self._recv(2)
-            length = struct.unpack("!H", l)[0]
+            length_data = self._recv(2)
+            length = struct.unpack("!H", length_data)[0]
         elif length == 0x7f:
-            l = self._recv(8)
-            length = struct.unpack("!Q", l)[0]
+            length_data = self._recv(8)
+            length = struct.unpack("!Q", length_data)[0]
 
+        mask_key = ""
         if mask:
             mask_key = self._recv(4)
         data = self._recv(length)
+        if traceEnabled:
+            recieved = header_bytes + length_data + mask_key + data
+            logger.debug("recv: " + repr(recieved))
+
         if mask:
             data = ABNF.mask(mask_key, data)