replay-chunk: accept capture as arg, py2 compat

The common case is replaying a capture, so avoid all tshark boilerplate
and run the command in the script. Fix the mixup of server and client
sockets ("if reply, then write from the server socket"). Allow
chunk size to be customized.

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)