From 8e86b9e90fa7bfc2274b1c801932cade5bccf564 Mon Sep 17 00:00:00 2001
From: Benjamin Poirier <benjamin.poirier@gmail.com>
Date: Thu, 27 Jul 2017 16:07:06 -0700
Subject: [PATCH] tcptop: Cleanup argument parsing

---
 tools/tcptop.py | 31 +++++++++++++++----------------
 1 file changed, 15 insertions(+), 16 deletions(-)

diff --git a/tools/tcptop.py b/tools/tcptop.py
index a6154c66..0f9ead6e 100755
--- a/tools/tcptop.py
+++ b/tools/tcptop.py
@@ -34,6 +34,13 @@ from subprocess import call
 import ctypes as ct
 
 # arguments
+def range_check(string):
+    value = int(string)
+    if value < 1:
+        msg = "value must be stricly positive, got %d" % (value,)
+        raise argparse.ArgumentTypeError(msg)
+    return value
+
 examples = """examples:
     ./tcptop           # trace TCP send/recv by host
     ./tcptop -C        # don't clear the screen
@@ -49,15 +56,11 @@ parser.add_argument("-S", "--nosummary", action="store_true",
     help="skip system summary line")
 parser.add_argument("-p", "--pid",
     help="trace this PID only")
-parser.add_argument("interval", nargs="?", default=1,
+parser.add_argument("interval", nargs="?", default=1, type=range_check,
     help="output interval, in seconds (default 1)")
-parser.add_argument("count", nargs="?", default=99999999,
+parser.add_argument("count", nargs="?", default=-1, type=range_check,
     help="number of outputs")
 args = parser.parse_args()
-countdown = int(args.count)
-if args.interval and int(args.interval) == 0:
-    print("ERROR: interval 0. Exiting.")
-    exit()
 debug = 0
 
 # linux stats
@@ -204,15 +207,13 @@ ipv6_recv_bytes = b["ipv6_recv_bytes"]
 print('Tracing... Output every %s secs. Hit Ctrl-C to end' % args.interval)
 
 # output
-exiting = 0
-while (1):
+i = 0
+exiting = False
+while i != args.count and not exiting:
     try:
-        if args.interval:
-            sleep(int(args.interval))
-        else:
-            sleep(99999999)
+        sleep(args.interval)
     except KeyboardInterrupt:
-        exiting = 1
+        exiting = True
 
     # header
     if args.noclear:
@@ -282,6 +283,4 @@ while (1):
     ipv6_send_bytes.clear()
     ipv6_recv_bytes.clear()
 
-    countdown -= 1
-    if exiting or countdown == 0:
-        exit()
+    i += 1
-- 
2.34.1