From 3c8d249cb5b3e9af4ffd9c0cad4df3d278bda1cd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fran=C3=A7ois=20De=20Keersmaeker?=
 <francois.dekeersmaeker@uclouvain.be>
Date: Thu, 11 May 2023 14:52:01 +0200
Subject: [PATCH] Updated -r option

---
 README.md           |  2 +-
 src/pcap-tweaker.py | 29 +++++++++++++++++++++++++----
 2 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index 437ae8a..6a8682b 100644
--- a/README.md
+++ b/README.md
@@ -45,7 +45,7 @@ It will be created if it doesn't exist.
 
 * `-h`, `--help`: show help message and exit
 * `-d`, `--dry-run`: don't write the output PCAP file (but still write the CSV log file)
-* `-r`, `--random-range`: upper bound for the random range, which will select for each packet if it will be edited or not. In practice, each packet will be edited with a probability of `1/(r+1)`. Default: `0` (edit all packets).
+* `-r`, `--random-range`: upper bound for the random range, which will select for each packet if it will be edited or not. In practice, each packet will be edited with a probability of `1/r`. Must be a strictly positive integer. Default: `1` (edit all packets).
 * `-n`, `--packet-number`: index of the packet to edit, starting from 1. Can be specified multiple times. If this is used, only the specified packets will be edited, and no random editing will be performed.
 
 
diff --git a/src/pcap-tweaker.py b/src/pcap-tweaker.py
index 1d3d5e6..75a9238 100644
--- a/src/pcap-tweaker.py
+++ b/src/pcap-tweaker.py
@@ -13,6 +13,24 @@ from scapy.contrib import coap, igmp, igmpv3
 from packet.Packet import Packet
 
 
+def strictly_positive_int(value: any) -> int:
+    """
+    Custom argparse type for a strictly positive integer value.
+    
+    :param value: argument value to check
+    :return: argument as integer if it is strictly positive
+    :raises argparse.ArgumentTypeError: if argument does not represent a strictly positive integer
+    """
+    try:
+        ivalue = int(value)
+    except ValueError:
+        raise argparse.ArgumentTypeError(f"{value} does not represent an integer.")
+    else:
+        if ivalue < 1:
+            raise argparse.ArgumentTypeError(f"{value} does not represent a strictly positive integer.")
+        return ivalue
+
+
 if __name__ == "__main__":
 
     # Script-related variables
@@ -32,11 +50,14 @@ if __name__ == "__main__":
     # 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.")
+    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=int, default=0, help="Upper bound for 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.")
+    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()
 
@@ -82,7 +103,7 @@ if __name__ == "__main__":
                 for packet in packets:
 
                     # Choose randomly if we edit this packet
-                    if random.randint(0, args.random_range) != 0:
+                    if random.randrange(0, args.random_range) != 0:
                         # Packet won't be edited
                         # Go to next packet
                         i += 1
-- 
GitLab