Merge.

darcs-hash:20090416055623-6342c-7aab4ded5a3ef751719907c602c36668fb2bb3f6.gz

utilities/tcptrace.py

@@ -37,6 +37,7 @@
 config=pyflag.conf.ConfObject()
 from pyflag.CacheManager import CachedWriter
 import pyflag.FlagFramework as FlagFramework
+import cStringIO
 
 parser = OptionParser(usage = """%prog [options] pcap_file ... pcap_file
 
@@ -64,6 +65,15 @@
 
 CONS = 0
 
+class MyCachedWriter(CachedWriter):
+    def __init__(self, filename):
+        self.filename = filename
+        self.fd = cStringIO.StringIO()
+        self.offset = 0
+
+    def close(self):
+        self.write_to_file()
+
 def Callback(mode, packet, connection, options = None):
     global CONS
     import socket, struct, time
@@ -76,8 +86,8 @@ def Callback(mode, packet, connection, options = None):
             connection['reverse']['con_id'] = CONS
             CONS +=1
 
-            connection['data'] = CachedWriter("%s/S%s" % (options.prefix, connection['con_id']))
-            connection['reverse']['data'] = CachedWriter("%s/S%s" % (options.prefix, connection['reverse']['con_id']))
+            connection['data'] = MyCachedWriter("%s/S%s" % (options.prefix, connection['con_id']))
+            connection['reverse']['data'] = MyCachedWriter("%s/S%s" % (options.prefix, connection['reverse']['con_id']))
 
         ip = packet.find_type("IP")
         connection['src_ip'] = ip.src