Refactored main

src/pcap-tweaker.py

@@ -31,42 +31,20 @@ def strictly_positive_int(value: any) -> int:
         return ivalue
 
 
-if __name__ == "__main__":
-
-    # Script-related variables
-    script_name = os.path.basename(__file__)
-    script_path = os.path.dirname(os.path.abspath(__file__))
-
-    ### LOGGING CONFIGURATION ###
-    logging.basicConfig(level=logging.INFO)
-    logging.info(f"Starting {script_name}")
-
-
-    ### ARGUMENT PARSING ###
-    parser = argparse.ArgumentParser(
-        prog=script_name,
-        description="Randomly edit packet fields in a PCAP file."
-    )
-    # Positional arguments: input PCAP file
-    parser.add_argument("input_pcaps", metavar="pcap", type=str, nargs="+", help="Input PCAP files.")
-    # Optional flag: -d / --dry-run
-    parser.add_argument("-d", "--dry-run", action="store_true",
-                        help="Dry run: do not write output PCAP file.")
-    # Optional flag: -r / --random-range
-    parser.add_argument("-r", "--random-range", type=strictly_positive_int, default=1,
-                        help="Upper bound for random range (not included). Must be a strictly positive integer. Default: 1 (edit each packet).")
-    # Optional flag: -n / --packet-number
-    parser.add_argument("-n", "--packet-number", type=int, action="append",
-                        help="Index of the packet to edit, starting form 1. Can be specifed multiple times.")
-    # Parse arguments
-    args = parser.parse_args()
-
-
-    ### MAIN PROGRAM ###
+def tweak_pcaps(pcaps: list, dry_run: bool, packet_numbers: list, random_range: int = 1) -> None:
+    """
+    Main functionality of the program:
+    (Randomly) edit packet fields in a (list of) PCAP file(s).
 
+    :param pcaps: list of input PCAP files
+    :param dry_run: if True, do not write output PCAP file
+    :param packet_numbers: list of packet numbers to edit (starting from 1)
+    :param random_range: upper bound for random range (not included)
+    """
+    
     # Loop on given input PCAP files
-    for input_pcap in args.input_pcaps:
-        # Useful paths
+    for input_pcap in pcaps:
+        # PCAP file directory
         input_dir = os.path.dirname(input_pcap)
 
         # Read input PCAP file
@@ -83,9 +61,9 @@ if __name__ == "__main__":
             writer = csv.DictWriter(csv_file, fieldnames=field_names)
             writer.writeheader()
 
-            if args.packet_number is not None:
+            if packet_numbers is not None:
                 # Edit specific packets
-                for i in args.packet_number:
+                for i in packet_numbers:
                     packet = packets[i - 1]  # -1 because packet numbers start at 1
                     try:
                         my_packet = Packet.init_packet(packet, i)
@@ -103,7 +81,7 @@ if __name__ == "__main__":
                 for packet in packets:
 
                     # Choose randomly if we edit this packet
-                    if random.randrange(0, args.random_range) != 0:
+                    if random.randrange(0, random_range) != 0:
                         # Packet won't be edited
                         # Go to next packet
                         i += 1
@@ -127,8 +105,42 @@ if __name__ == "__main__":
         os.makedirs(output_dir, exist_ok=True)
         output_pcap = os.path.basename(input_pcap).replace(".pcap", ".edit.pcap")
         output_pcap = os.path.join(output_dir, output_pcap)
-        if args.dry_run:
+        if dry_run:
             logging.info(f"Dry run: did not write output PCAP file: {output_pcap}")
         else:
             scapy.wrpcap(output_pcap, packets)
             logging.info(f"Wrote output PCAP file: {output_pcap}")
+
+
+if __name__ == "__main__":
+
+    # This script's name
+    script_name = os.path.basename(__file__)
+
+    ### LOGGING CONFIGURATION ###
+    logging.basicConfig(level=logging.INFO)
+    logging.info(f"Starting {script_name}")
+
+
+    ### ARGUMENT PARSING ###
+    parser = argparse.ArgumentParser(
+        prog=script_name,
+        description="Randomly edit packet fields in a PCAP file."
+    )
+    # Positional arguments: input PCAP file
+    parser.add_argument("input_pcaps", metavar="pcap", type=str, nargs="+", help="Input PCAP files.")
+    # Optional flag: -d / --dry-run
+    parser.add_argument("-d", "--dry-run", action="store_true",
+                        help="Dry run: do not write output PCAP file.")
+    # Optional flag: -r / --random-range
+    parser.add_argument("-r", "--random-range", type=strictly_positive_int, default=1,
+                        help="Upper bound for random range (not included). Must be a strictly positive integer. Default: 1 (edit each packet).")
+    # Optional flag: -n / --packet-number
+    parser.add_argument("-n", "--packet-number", type=int, action="append",
+                        help="Index of the packet to edit, starting form 1. Can be specifed multiple times.")
+    # Parse arguments
+    args = parser.parse_args()
+
+
+    ### MAIN PROGRAM ###
+    tweak_pcaps(args.input_pcaps, args.dry_run, args.packet_number, args.random_range)