diff --git a/src/pcap-tweaker.py b/src/pcap-tweaker.py index 75a92382c2cc32bdddb9eb6f4baf679a3945dc0c..a52a834930a2ca8492a6432f7a0bfe1dd0e20a57 100644 --- a/src/pcap-tweaker.py +++ b/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)