add csv export mode

.gitignore

@@ -4,6 +4,7 @@
 ### Project dependent ###
 /capture.pcap
 .env
+results.csv
 
 ### Linux ###
 *~

analyse_tcp_packets.py

@@ -4,19 +4,7 @@
 import socket
 import os
 import argparse
-
-parser = argparse.ArgumentParser()
-parser.add_argument('--captureFileName', '-cfn',
-                    help='The name of the pcap file to analyse, default: capture.pcap')
-args = parser.parse_args()
-
-LISTENED_IP = os.getenv('LISTENED_IP')
-captureFileName = "capture.pcap" if args.captureFileName is None else args.captureFileName
-
-
-# Check if capture file exists
-if not os.path.exists(captureFileName):
-    raise Exception("File {} doesn't exist".format(captureFileName))
+import csv_saver
 
 
 # Analyse a pkt to save it in the good key of our date structure
@@ -57,6 +45,7 @@ def analyse_packets(pkt):
 def get_packet_size(pkt):
     return int(pkt.length.raw_value, 16) * 0.000001
 
+
 # Save a new stream and its first packet in the dict
 def save_new_stream(stream_id, timestamp, ip, pkt):
     domain = reverse_dns(ip)
@@ -71,10 +60,15 @@ def save_new_stream(stream_id, timestamp, ip, pkt):
         'endTime': timestamp,
     }
 
+
 # Send a group of packets that seems to be together to the DB
 def push_data(key):
     print('Push data: ' + str(packet_dict[key]))
-    db.save_element(packet_dict[key], captureFileName)
+    if (args.export == "mongo"):
+        db.save_element(packet_dict[key], captureFileName)
+    else:
+        csv_saver.save_element(packet_dict[key])
+
 
 # Reverse DNS a remote IP
 def reverse_dns(ip):
@@ -85,6 +79,21 @@ def reverse_dns(ip):
         return ""
 
 
+# Arguments available
+parser = argparse.ArgumentParser()
+parser.add_argument('--captureFileName', '-cfn',
+                    help='The name of the pcap file to analyse, default: capture.pcap')
+parser.add_argument('--export', '-e',
+                    help="The export mode you want to use: mongo or csv, default: csv")
+args = parser.parse_args()
+
+LISTENED_IP = os.getenv('LISTENED_IP')
+
+captureFileName = "capture.pcap" if args.captureFileName is None else args.captureFileName
+# Check if capture file exists
+if not os.path.exists(captureFileName):
+    raise Exception("File {} doesn't exist".format(captureFileName))
+
 # Data structure
 """ packet_dict = {
   stream_index: {
@@ -100,8 +109,10 @@ def reverse_dns(ip):
 } """
 packet_dict = {}
 
-# Connect to db
-db.connect_to_db()
+# Connect to db if we are in the mongo export mode
+if (args.export == "mongo"):
+    environment.check_mongo_env()
+    db.connect_to_db()
 
 # Open the capture file
 cap = pyshark.FileCapture(captureFileName)

csv_saver.py

@@ -0,0 +1,17 @@
+import csv
+import os
+
+csv_columns = ['sumDelta', 'averageDelta', 'ip', 'domain',
+               'numberOfPackets', 'totalMbSize', 'startTime', 'endTime']
+
+# Insert a row in the results csv file
+def save_element(element):
+    if not os.path.exists('results.csv'):
+        with open('results.csv', 'w') as csvfile:
+            writer = csv.DictWriter(csvfile, fieldnames=csv_columns)
+            writer.writeheader()
+            writer.writerow(element)
+    else:
+        with open('results.csv', 'a') as csvfile:
+            writer = csv.DictWriter(csvfile, fieldnames=csv_columns)
+            writer.writerow(element)

db.py

@@ -14,7 +14,7 @@ def connect_to_db():
         "mongodb+srv://{}:{}@{}/test?retryWrites=true&w=majority".format(MONGO_DB_USER, MONGO_DB_PASSWORD, MONGO_CLUSTER_ADDRESS))
     db = cluster[MONGO_DB_NAME]
 
-# Insert a document in our analysed_packets collection
+# Insert a document in the given collection
 def save_element(element, collection_name):
     global db
     if db is None:

environment.py

@@ -2,12 +2,20 @@
 from dotenv import load_dotenv
 load_dotenv()
 
-if (os.getenv('MONGO_CLUSTER_ADDRESS') is None or os.getenv('MONGO_CLUSTER_ADDRESS') == ""
-    or os.getenv('MONGO_DB_NAME') is None or os.getenv('MONGO_DB_NAME') == ""
-    or os.getenv('MONGO_DB_USER') is None or os.getenv('MONGO_DB_USER') == ""
-    or os.getenv('MONGO_DB_PASSWORD') is None or os.getenv('MONGO_DB_PASSWORD') == ""
-    or os.getenv('LISTENED_IP') is None or os.getenv('LISTENED_IP') == ""
-    or os.getenv('LOCAL_IP') is None or os.getenv('LOCAL_IP') == ""
-    or os.getenv('INTERFACE') is None or os.getenv('INTERFACE') == ""):
-    raise Exception(
-        'Please complete all the environment variables in the .env file')
+# Check if we have all the environment variables for the sniffer
+def check_sniffer_env():
+    if (os.getenv('LISTENED_IP') is None or os.getenv('LISTENED_IP') == ""
+        or os.getenv('LOCAL_IP') is None or os.getenv('LOCAL_IP') == ""
+        or os.getenv('INTERFACE') is None or os.getenv('INTERFACE') == ""):
+        raise Exception(
+            '\n\nPlease complete the following environment variables in the .env file:\nLISTENED_IP\nLOCAL_IP\nINTERFACE\n')
+
+
+# Check if we have all the environment variables for mongoDB
+def check_mongo_env():
+    if (os.getenv('MONGO_CLUSTER_ADDRESS') is None or os.getenv('MONGO_CLUSTER_ADDRESS') == ""
+        or os.getenv('MONGO_DB_NAME') is None or os.getenv('MONGO_DB_NAME') == ""
+        or os.getenv('MONGO_DB_USER') is None or os.getenv('MONGO_DB_USER') == ""
+        or os.getenv('MONGO_DB_PASSWORD') is None or os.getenv('MONGO_DB_PASSWORD') == ""):
+        raise Exception(
+            '\n\nPlease complete the following environment variables in the .env file:\nMONGO_CLUSTER_ADDRESS\nMONGO_DB_NAME\nMONGO_DB_USER\nMONGO_DB_PASSWORD\n')

sniffer.py

@@ -3,6 +3,10 @@
 import os
 import argparse
 
+# Check if we have the necessaries environment variables
+environment.check_sniffer_env()
+
+# Arguments available
 parser = argparse.ArgumentParser()
 parser.add_argument('--outputFile', '-of',
                     help='The name of the pcap output file, default: capture')
@@ -20,6 +24,7 @@
 capture = pyshark.LiveCapture(
     interface=INTERFACE, output_file=output_file, display_filter=filter)
 
+# Launch the capture
 if args.timeout is None:
     capture.sniff()
 else: