1 files changed, 45 insertions, 10 deletions

diff --git a/crafted-pkt/replay-chunks.py b/crafted-pkt/replay-chunks.py
index 404e2ed..b969b27 100755
--- a/crafted-pkt/replay-chunks.py
+++ b/crafted-pkt/replay-chunks.py
@@ -1,16 +1,35 @@
 #!/usr/bin/env python
 #
-# Replay a communication, splitting data at a certain chunk size
-# (hard-coded to 2). This can be used to test reassembly for example.
+# Replay a communication, splitting data at a certain chunk sizes
+# (defaults to 2). This can be used to test reassembly for example.
 #
 # Copyright (C) 2014 Peter Wu <peter@lekensteyn.nl>
 
 # Usage (assuming a capture file with TCP stream 0 at loopback interface lo)
-# dumpcap -i lo -w split.pcapng
-# tshark -r old-capture.pcapng -q -z follow,tcp,raw,0 | ./replay-chunks.py
+#     dumpcap -i lo -w split.pcapng
+#     ./replay-chunks.py old-capture.pcapng
+# Run ./replay-chunks.py --help for details
 
 import socket
 import sys
+import os
+from argparse import ArgumentParser
+from subprocess import PIPE
+from subprocess import Popen as _Popen
+if hasattr(_Popen, '__exit__'):
+    Popen = _Popen
+else:
+    class PopenClosable(_Popen):
+        def __enter__(self):
+            return self
+        def __exit__(self, exc_type, exc_value, traceback):
+            self.stdout.close()
+    Popen = PopenClosable
+try:
+    from subprocess import DEVNULL
+except ImportError:
+    # Python < 3.3 compatibility
+    DEVNULL = open(os.devnull, 'w')
 
 state = 'init'
 
@@ -32,7 +51,7 @@ class FollowParser(object):
         self.chunk_size = chunk_size
 
     def add_data(self, data, is_reply):
-        sock = self.sock_client if is_reply else self.sock_server
+        sock = self.sock_server if is_reply else self.sock_client
         for i in range(0, len(data), self.chunk_size):
             sock.send(data[i:i+self.chunk_size])
         print('{}: {}'.format('S->C' if is_reply else 'C->S', _dumpbytes(data)))
@@ -60,7 +79,8 @@ class FollowParser(object):
         pass
 
     def open_sockets(self):
-        with socket.socket() as svr:
+        svr = socket.socket()
+        try:
             svr.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
             svr.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
             svr.bind(self.addr)
@@ -68,6 +88,8 @@ class FollowParser(object):
             self.sock_client = socket.socket()
             self.sock_client.connect(self.addr)
             self.sock_server, remote_addr = svr.accept()
+        except:
+            svr.close()
 
     def feed_data(self, line):
         old_state = self.state
@@ -83,15 +105,28 @@ class FollowParser(object):
             self.sock_client.close()
             self.sock_client = None
 
-def main():
-    parser = FollowParser()
+def main(tshark_output, chunk_size):
+    parser = FollowParser(chunk_size=chunk_size)
     try:
-        for line in sys.stdin:
+        for line in tshark_output:
             old_state, new_state = parser.feed_data(line)
             if new_state != old_state and old_state == parser.state_find_node_1:
                 print('Found server node: {}:{}'.format(*parser.addr))
     finally:
         parser.close()
 
+parser = ArgumentParser(description='Replay TCP capture')
+parser.add_argument('-s', '--chunk-size', type=int, default=2,
+                    help='Maximum size of each chunk (default %(default)d)')
+parser.add_argument('file', help='Any capture format recognized by tshark')
 if __name__ == '__main__':
-    main()
+    _args = parser.parse_args()
+    _cmd = [
+        'tshark',
+        '-r',
+        _args.file,
+        '-q',
+        '-z', 'follow,tcp,raw,0'
+    ]
+    with Popen(_cmd, stdin=DEVNULL, stdout=PIPE, universal_newlines=True) as p:
+        main(p.stdout, chunk_size=_args.chunk_size)