Merge pull request #222 from anarkiwi/annotate2

Annotate pcaps with protocol by default.

Author: cglewis

pcap_to_node_pcap/pcap_to_node_pcap.py

@@ -6,6 +6,8 @@
 @author: Blake Pagon
 """
 
+# TODO: https://github.com/PyCQA/bandit/issues/333 for bandit false positive on subprocess.
+
 import argparse
 import datetime
 import ipaddress
@@ -27,13 +29,26 @@ def ipaddress_fields(json_fields):
         ipas.add(str(ipa))
     return ipas
 
+def pcap_name_with_layers(pcap_filename, pcap_layers, pcap_suffix):
+    pcap_basename = os.path.basename(pcap_filename)
+    pcap_basename = pcap_basename.replace(pcap_suffix, '')
+    layers_str = '-'.join(pcap_layers)
+    layers_pcap_filename = pcap_filename.replace(
+        pcap_basename, '-'.join((pcap_basename, layers_str)))
+    return layers_pcap_filename
+
 def proto_annotate_pcaps(pcap_dir):
-    pap_filenames = [
-        pcap.path for pcap in os.scandir(pcap_dir)
-        if pcap.is_file() and pcap.path.endswith('pcap')]
+    pcap_suffix = '.pcap'
+    try:
+        pap_filenames = [
+            pcap.path for pcap in os.scandir(pcap_dir)
+            if pcap.is_file() and pcap.path.endswith(pcap_suffix)]
+    except FileNotFoundError as err:
+        print(err)
+        return
     for pcap_filename in pap_filenames:
         try:
-            response = subprocess.check_output(shlex.split(' '.join(
+            response = subprocess.check_output(shlex.split(' '.join( # nosec
                     ['./tshark', '-T', 'json', '-c', str(10), '-r', pcap_filename])))
             pcap_json = json.loads(response.decode("utf-8"))
         except (json.decoder.JSONDecodeError, subprocess.CalledProcessError) as e:
@@ -52,10 +67,7 @@ def proto_annotate_pcaps(pcap_dir):
             packet_layers = list(ipas) + list(layers_json.keys())
             if len(packet_layers) > len(pcap_layers):
                 pcap_layers = packet_layers
-        pcap_basename = os.path.basename(pcap_filename)
-        layers_str = '-'.join(pcap_layers)
-        layers_pcap_filename = pcap_filename.replace(
-            pcap_basename, '-'.join((layers_str, pcap_basename)))
+        layers_pcap_filename = pcap_name_with_layers(pcap_filename, pcap_layers, pcap_suffix)
         os.rename(pcap_filename, layers_pcap_filename)
 
 def connect_rabbit(host='messenger', port=5672, queue='task_queue'):
@@ -98,43 +110,42 @@ def run_tool(path, protoannotate):
 
     # need to make directories to store results from pcapsplitter
     base_dir = path.rsplit('/', 1)[0]
-    timestamp = ""
-    try:
-        timestamp = '-'.join(str(datetime.datetime.now()).split(' ')) + '-UTC'
-        timestamp = timestamp.replace(':', '_')
-    except Exception as e:  # pragma: no cover
-        print("couldn't create output directory with unique timestamp")
+    timestamp = '-'.join(str(datetime.datetime.now()).split(' ')) + '-UTC'
+    timestamp = timestamp.replace(':', '_')
     # make directory for tool name recognition of piping to other tools
     output_dir = os.path.join(base_dir, 'pcap-node-splitter' + '-' + timestamp)
-    try:
-        os.mkdir(output_dir)
-        os.mkdir(output_dir + '/clients')
-        os.mkdir(output_dir + '/servers')
-    except OSError:  # pragma: no cover
-        print("couldn't make directories for output of this tool")
     clients_dir = os.path.join(output_dir, 'clients')
     servers_dir = os.path.join(output_dir, 'servers')
+    for new_dir in (output_dir, clients_dir, servers_dir):
+        try:
+            os.mkdir(new_dir)
+        except OSError as err:
+            print("couldn't make directory %s for output of this tool: %s" % (new_dir, err))
 
-    try:
-        subprocess.check_call(shlex.split("./PcapSplitter -f " +
-                                          path + " -o " + clients_dir + " -m client-ip"))
-
-        subprocess.check_call(shlex.split("./PcapSplitter -f " +
-                                          path + " -o " + servers_dir + " -m server-ip"))
-    except Exception as e:
-        print(str(e))
+    for tool_cmd in (
+            " ".join(("./PcapSplitter -f", path, "-o", clients_dir, "-m client-ip")),
+            " ".join(("./PcapSplitter -f", path, "-o", servers_dir, "-m server-ip"))):
+        try:
+            subprocess.check_call(shlex.split(tool_cmd)) # nosec
+        except Exception as err:
+            print("%s: %s" % (tool_cmd, err))
 
     if protoannotate:
         for pcap_dir in (clients_dir, servers_dir):
             proto_annotate_pcaps(pcap_dir)
 
-    return output_dir + '/clients'
+    return clients_dir
 
-if __name__ == '__main__':  # pragma: no cover
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--protoannotate', help='use tshark to annotate pcaps with protocol', action='store_true')
+def parse_args(parser):
+    parser.add_argument('--protoannotate', help='use tshark to annotate pcaps with protocol',
+        action='store_true', default=True)
     parser.add_argument('paths', nargs='*')
     args = parser.parse_args()
+    return args
+
+
+if __name__ == '__main__':  # pragma: no cover
+    args = parse_args(argparse.ArgumentParser())
     path = get_path(args.paths)
     if path:
         result_path = run_tool(path, args.protoannotate)

pcap_to_node_pcap/test_pcap_to_node_pcap.py

@@ -4,10 +4,22 @@
 Created on 20 December 2018
 @author: Charlie Lewis
 """
+import argparse
 import sys
 
-from .pcap_to_node_pcap import get_path
-from .pcap_to_node_pcap import run_tool
+from .pcap_to_node_pcap import get_path, run_tool, pcap_name_with_layers, parse_args
+
+
+def test_pcap_name_with_layers():
+    pcap_suffix = '.pcap'
+    pcap_basename = 'trace_3cf8009a09d9684250ffa77d6f7752aee61463a8_2019-11-07_04_11_19-server-ip-74-125-68-189'
+    pcap_filename = pcap_basename + pcap_suffix
+    new_name = pcap_name_with_layers(pcap_filename, ['a', 'b', 'c'], pcap_suffix)
+    assert new_name == pcap_basename + '-a-b-c' + pcap_suffix
+
+def test_parse_args():
+    args = parse_args(argparse.ArgumentParser([]))
+    assert args.protoannotate
 
 
 def test_get_path():