support spice1.proto, output surface id's in right column, small fixes

5 files changed, 81 insertions, 40 deletions

diff --git a/client.py b/client.py
index aa164cb..9311f06 100644
--- a/client.py
+++ b/client.py
@@ -53,7 +53,7 @@ class SpiceDataHeader(Struct):
             logger.error('bad type of packet - %s' % inst)
             import pdb; pdb.set_trace()
             return None
-        if inst.e.size > 1000000:
+        if inst.e.size > 2000000:
             logger.error('too large packet - %s' % inst)
             import pdb; pdb.set_trace()
             return None
@@ -89,6 +89,17 @@ class SpiceLinkHeader(Struct):
             return None
         return inst
 
+    @classmethod
+    def set_proto(cls, *args, **kw):
+        inst = cls(*args, **kw)
+        if inst.e.magic != SPICE_MAGIC:
+            logger.error('bad magic in packet: %s' % inst)
+            return None
+        import client_proto
+        client_proto.set_proto(major_version=inst.e.major_version,
+            minor_version=inst.e.minor_version)
+        return inst
+
 link_header_size = SpiceLinkHeader.size
 
 class SpiceLinkMess(Struct):
diff --git a/client_proto.py b/client_proto.py
index 3095243..a90f0fa 100644
--- a/client_proto.py
+++ b/client_proto.py
@@ -16,22 +16,28 @@ if not 'proto' in locals():
     sys.path.append('../spice/python_modules')
     import spice_parser
     import ptypes
-    proto = spice_parser.parse('../spice/spice.proto')
+    proto = None
+    major_version, minor_version = None, None
+    #channels, num_spice_messages, valid_spice_messages
+    #all_channels
     reloads = 1
 else:
     reloads += 1
 
 
-channels = {}
-for channel in proto.channels:
-    channels[channel.value] = channel
-    channel.client = dict(zip([x.value for x in
-        channel.channel_type.client_messages],
-        channel.channel_type.client_messages))
-    channel.server = dict(zip([x.value for x in
-        channel.channel_type.server_messages],
-        channel.channel_type.server_messages))
-    # todo parsing of messages / building of messages
+def make_channels_dict(proto):
+    channels = {}
+    for channel in proto.channels:
+        channels[channel.value] = channel
+        channel.client = dict(zip([x.value for x in
+            channel.channel_type.client_messages],
+            channel.channel_type.client_messages))
+        channel.server = dict(zip([x.value for x in
+            channel.channel_type.server_messages],
+            channel.channel_type.server_messages))
+        # todo parsing of messages / building of messages
+    return channels
+channels = None
 
 def mapdict(f, d):
     return dict((k, f(v)) for k,v in d.items())
@@ -333,3 +339,29 @@ class NoPrint(object):
     def __len__(self):
         return len(self.s)
 
+def set_proto(major_version, minor_version):
+    global proto, channels, num_spice_messages, valid_spice_messages
+    global all_channels, valid_message_ids
+    if globals()['major_version'] == major_version and globals()['minor_version'] == minor_version:
+        return
+    globals().update(dict(major_version=major_version, minor_version=minor_version))
+    if major_version == 1 :
+        proto = spice_parser.parse('../spice/spice1.proto')
+    else:
+        proto = spice_parser.parse('../spice/spice.proto')
+    channels = make_channels_dict(proto)
+
+    num_spice_messages = sum(len(ch.client) + len(ch.server) for ch in
+        channels.values())
+
+    valid_message_ids = set(sum([ch.client.keys() + ch.server.keys() for ch in
+        channels.values()], []))
+
+    all_channels = channels.keys()
+
+def possible_channels(server_message, header):
+    return set(c for c in all_channels if header.e.type in
+        (client_proto.channels[c].server.keys() if server_message
+         else client_proto.channels[c].client.keys()))
+
+
diff --git a/pcapspice.py b/pcapspice.py
index 29c7027..a8f335a 100644
--- a/pcapspice.py
+++ b/pcapspice.py
@@ -10,25 +10,12 @@ logger = logging.getLogger('pcapspice')
 
 def is_single_packet_data(payload):
     return len(payload) == SpiceDataHeader(payload).e.size + SpiceDataHeader.size
-
-num_spice_messages = sum(len(ch.client) + len(ch.server) for ch in
-    client_proto.channels.values())
-
-valid_message_ids = set(sum([ch.client.keys() + ch.server.keys() for ch in
-    client_proto.channels.values()], []))
-
-all_channels = client_proto.channels.keys()
-
-def possible_channels(server_message, header):
-    return set(c for c in all_channels if header.e.type in 
-        (client_proto.channels[c].server.keys() if server_message
-         else client_proto.channels[c].client.keys()))
-
+#
 guesses = {}
 
 def guess_channel_iter():
     channel = None
-    optional_channels = set(all_channels)
+    optional_channels = set(client_proto.all_channels)
     seen_headers = []
     while True:
         src, dst, data_header = yield channel
@@ -36,7 +23,7 @@ def guess_channel_iter():
         if guesses.has_key(key):
             optional_channels = set([guesses[key]])
         seen_headers.append((src, dst, data_header))
-        optional_channels = optional_channels & possible_channels(
+        optional_channels = optional_channels & client_proto.possible_channels(
             src == min(src, dst), data_header)
         logger.debug(str(optional_channels))
         if len(optional_channels) == 1:
@@ -68,7 +55,7 @@ def channel_spice_message_iter(src, dst, guesser):
                     src, dst, data_header.e.serial, serial[src]))
                 bad_header = True
         serial[src] = data_header.e.serial
-        if data_header.e.type not in valid_message_ids:
+        if data_header.e.type not in client_proto.valid_message_ids:
             logger.error("bad message type %s in %s->%s" % (data_header.e.type,
                 src, dst))
             bad_header = True
@@ -85,13 +72,15 @@ def channel_spice_message_iter(src, dst, guesser):
 class ChannelGuesser(object):
 
     def __init__(self):
-        self.iter = guess_channel_iter()
-        self.iter.next()
+        self.iter = None
         self.channel = None
 
     def send(self, (src, dst, payload)):
         if self.channel:
             return self.channel
+        if self.iter is None:
+            self.iter = guess_channel_iter()
+            self.iter.next()
         self.channel = self.iter.send((src, dst, payload))
         return self.channel
 
@@ -112,7 +101,7 @@ def spice_iter(packet_iter):
         logger.debug('%s %s **** make_start_iter %s -> %s' %( guesser, is_server, src, dst))
         link_messages = []
         yield (SpiceLinkHeader.size,
-            SpiceLinkHeader.verify,
+            SpiceLinkHeader.set_proto,
             lambda h: h.e.size,
             (guesser.on_server_link_message if is_server else 
             guesser.on_client_link_message))
@@ -132,12 +121,14 @@ def spice_iter(packet_iter):
             lambda h: h.e.size,
             lambda pkt, src, dst, header:
                     message_iter.send((src, dst, (header, pkt)))))
+    def filter_result(r):
+        return (r.msg is not None and r.msg.data != None
+                    and hasattr(r.msg.data, 'result_name'))
     return header_conversation_iter(
         packet_iter,
         server_start=make_start_iter,
         client_start=make_start_iter,
         server=make_iter,
         client=make_iter,
-        filter_result=lambda r: (r.msg is not None and r.msg.data != None
-            and hasattr(r.msg.data, 'result_name')))
+        filter_result=filter_result)
 
diff --git a/pcaputil.py b/pcaputil.py
index c6432fe..048e156 100644
--- a/pcaputil.py
+++ b/pcaputil.py
@@ -242,6 +242,7 @@ def collect_packets(header_iter_gen):
     def collector(start_src, start_dst):
         packets = []
         history = []
+        output = []
         src, dst = start_src, start_dst
         header_iter = header_iter_gen(src, dst)
         hdr_size, hdr_ctor, size_from_hdr, pkt_ctor = header_iter.next()
@@ -249,7 +250,9 @@ def collect_packets(header_iter_gen):
         header = pkt = None
         searched_size = hdr_size
         while True:
-            src, dst, payload = yield CollectorResult(src, dst, msg)
+            collector_result = CollectorResult(src, dst, msg)
+            output.append(collector_result)
+            src, dst, payload = yield collector_result
             packets.append(payload)
             history.append(payload)
             logger.debug("collect_packets: (%s->%s) searching for %s (%s)" % (src, dst,
diff --git a/spicedump.py b/spicedump.py
index a37b2f6..3fcab25 100755
--- a/spicedump.py
+++ b/spicedump.py
@@ -1,14 +1,17 @@
 #!/usr/bin/env python
 import sys
-import pcaputil
-from proxy import proxy, closeallsockets
 from collections import defaultdict
 from itertools import izip_longest
 from time import time
 from select import select
 from optparse import OptionParser
+from itertools import chain,repeat
+import textwrap
 import logging
 
+import pcaputil
+from proxy import proxy, closeallsockets
+
 dt = 1.0
 verbose = 0
 
@@ -47,7 +50,8 @@ class SurfaceStatistics(object):
         for k, v in per_surface.items():
             if k == 'stream_create' and len(v) > 1:
                 import pdb; pdb.set_trace()
-            yield '%20s: %s' % (k, ','.join(map(str, v)))
+            for i, l in enumerate(textwrap.wrap(','.join(map(str, v)), width=40)):
+                yield '%20s: %s' % (k if i == 0 else '', l)
 
 def spicedump(p, stdscr=None):
     import pcapspice
@@ -55,7 +59,7 @@ def spicedump(p, stdscr=None):
     hist = Histogram()
     surface_stat = SurfaceStatistics()
     last_print = start_time = time()
-    messages = ['welcome to spicedump (spice quest?)']
+    messages = []
     if stdscr:
         stdscr.erase()
     # replace the "for d in spice:" loop with a select
@@ -69,7 +73,7 @@ def spicedump(p, stdscr=None):
             result_value = d.msg.data.result_value
             if (any(x in result_name for x in ['surface', 'stream'])
                 and not result_name in
-                    ['stream_data', 'stream_clip', 'stream_destroy']):
+                    ['stream_data', 'stream_clip', 'stream_destroy', 'stream_destroy_all']):
                 msg_d = dict(result_value)
                 if 'surface_id' in msg_d:
                     surface_stat.add(msg_d['surface_id'], result_name)