Unit tests for entire package

pcap_anonymize/__init__.py

-from .pcap_anonymize import anonymize_pcap
+from .pcap_anonymize import anonymize_pcap, anonymize_pcaps_in_dir

pcap_anonymize/__main__.py

-import os
 import argparse
-from .pcap_anonymize import anonymize_pcap
+from .pcap_anonymize import anonymize_pcap, anonymize_pcaps_in_dir
 
 
 ### MAIN FUNCTION ###
 def main() -> None:
     parser = argparse.ArgumentParser(description="Anonymize a PCAP traffic capture.")
-    parser.add_argument("input", type=str, help="Path to the input PCAP file.")
+    parser.add_argument("-i", "--input", type=str, help="Path to the input PCAP file.")
     parser.add_argument("-o", "--output", type=str, help="Path to the output PCAP file.")
+    parser.add_argument("-d", "--dir", type=str, help="Path to the directory containing the input PCAP files.")
     args = parser.parse_args()
 
-    anonymize_pcap(args.input, args.output)
+    if args.dir:
+        anonymize_pcaps_in_dir(args.dir)
+    else:
+        anonymize_pcap(args.input, args.output)
 
 
 ### ENTRY POINT ###

pcap_anonymize/pcap_anonymize.py

@@ -3,6 +3,7 @@ Anonymize all packets in a PCAP file.
 """
 
 import os
+import glob
 from pathlib import Path
 from scapy.all import Packet, sniff, wrpcap
 # Packet layers
@@ -69,13 +70,34 @@ def anonymize_pcap(input: os.PathLike, output: os.PathLike = None) -> None:
     Args:
         input: path to the input PCAP file
         output: path to the output PCAP file.
-                If None, create a new file having the same name as the input file with the suffix '.anonymized.pcap'.
+                If None, create a new file having the same name as the input file with the suffix '.anon.pcap'.
     """
+    global packets
+
     if output is None:
-        output = str(Path(input).with_suffix('.anonymized.pcap'))
+        output = str(Path(input).with_suffix(".anon.pcap"))
 
     # Read and anonymize packets from the input file
     sniff(offline=input, prn=anonymize_packet, store=False)
 
     # Write anonymized packets to the output file
     wrpcap(output, packets)
+
+    # Reset global packets list
+    packets = []
+
+
+def anonymize_pcaps_in_dir(dir: os.PathLike) -> None:
+    """
+    Anonymize all PCAP files in a directory.
+
+    Args:
+        dir: path to the directory containing the PCAP files
+    """
+    for pcap_file in glob.glob(os.path.join(dir, "*.pcap")):
+
+        # Skip traces already anonymized
+        if pcap_file.endswith(".anon.pcap"):
+            continue
+
+        anonymize_pcap(pcap_file, None)

test/test_traces.py

+"""
+Test the package with PCAP traces.
+"""
+
+import os
+import glob
+from pathlib import Path
+from pcap_anonymize import anonymize_pcap, anonymize_pcaps_in_dir
+
+
+### TEST CONSTANTS ###
+
+dir_self = os.path.dirname(os.path.abspath(__file__))
+dir_traces = os.path.join(dir_self, "traces")
+
+
+### TEST FUNCTIONS ###
+
+def test_anonymize_pcap_http(tmp_path: Path) -> None:
+    """
+    Test the package with a PCAP trace containing HTTP packets.
+
+    Args:
+        tmp_path (pathlib.Path): temporary directory to store the output traces
+    """
+    input = os.path.join(dir_traces, "http.pcap")
+    output = os.path.join(tmp_path, "http.anon.pcap")
+    anonymize_pcap(input, output)
+    assert os.path.exists(input)
+    assert os.path.exists(output)
+
+
+def test_anonymize_pcap_dhcp(tmp_path: Path) -> None:
+    """
+    Test the package with a PCAP trace containing DHCP packets.
+
+    Args:
+        tmp_path (pathlib.Path): temporary directory to store the output traces
+    """
+    input = os.path.join(dir_traces, "dhcp.pcap")
+    output = os.path.join(tmp_path, "dhcp.anon.pcap")
+    anonymize_pcap(input, output)
+    assert os.path.exists(input)
+    assert os.path.exists(output)
+
+
+def test_anonymize_pcap_tplink(tmp_path: Path) -> None:
+    """
+    Test the package with a PCAP trace containing
+    TP-Link Smart Home protocol packets.
+
+    Args:
+        tmp_path (pathlib.Path): temporary directory to store the output traces
+    """
+    input = os.path.join(dir_traces, "tplink.pcap")
+    output = os.path.join(tmp_path, "tplink.anon.pcap")
+    anonymize_pcap(input, output)
+    assert os.path.exists(input)
+    assert os.path.exists(output)
+
+
+def test_anonymize_pcaps_in_dir() -> None:
+    """
+    Test the package with a directory containing PCAP traces.
+    """
+    input_traces = glob.glob(os.path.join(dir_traces, "*.pcap"))
+    anonymize_pcaps_in_dir(dir_traces)
+    output_traces = glob.glob(os.path.join(dir_traces, "*.anon.pcap"))
+    assert len(output_traces) == len(input_traces)
+
+    # Clean up
+    for output_trace in output_traces:
+        os.remove(output_trace)

test/traces/.gitignore

+# Anonymized PCAP files
+*.anon.pcap