Rewrite pingnet.py to use multiple threads

--HG--
extra : amend_source : 595a48f6b8cd4de9bc22e5731e816240e63ab8ff

pingnet.py

@@ -1,36 +1,113 @@
-#!/usr/bin/env python3
+#!/usr/bin/env python
+import argparse
 import ipaddr
 import os
-import socket
+import signal
 import subprocess
 import sys
+import threading
+import time
+try:
+    import queue
+except ImportError: # Python 2
+    import Queue as queue
 
-def main():
-    for arg in sys.argv[1:]:
-        try:
-            net = ipaddr.IPv4Network(arg)
-        except ipaddr.AddressValueError:
-            sys.stderr.write('Invalid network address: {}\n'.format(arg))
+
+class WaitQueue(queue.Queue):
+
+    def __init__(self, *args, **kwargs):
+        self.event =  threading.Event()
+        queue.Queue.__init__(self, *args, **kwargs)
+
+
+class Queuer(threading.Thread):
+
+    def __init__(self, iterable):
+        self.iterable = iterable
+        self.queue = WaitQueue()
+        self.go = True
+        super(Queuer, self).__init__()
+
+    def stop(self):
+        self.go = False
+
+    def run(self):
+        for item in self.iterable:
+            if not self.go:
+                break
+            self.queue.put(item)
+            self.queue.event.set()
+
+
+class Pinger(threading.Thread):
+
+    def __init__(self, q):
+        self.queue = q
+        self.go = True
+        self.quiet = False
+        super(Pinger, self).__init__()
+
+    def stop(self):
+        self.go = False
+
+    def run(self):
+        self.queue.event.wait()
+        while self.go:
+            try:
+                ip = self.queue.get(False)
+            except queue.Empty:
+                break
+
+            try:
+                self.ping(str(ip))
+            except KeyboardInterrupt:
+                break
+
+    def ping(self, host):
+        p = subprocess.Popen(['ping', '-c', '1', '-W', '1', host],
+                             stdin=open(os.devnull),
+                             stdout=open(os.devnull, 'w'),
+                             stderr=subprocess.STDOUT)
+        p.wait()
+        if p.returncode == 0:
+            if not self.quiet:
+                sys.stdout.write('{0} is in use\n'.format(host))
         else:
-            for ip in net.iterhosts():
-                p = subprocess.Popen(['ping', '-c1', '-W1', str(ip)],
-                                     stdout=open(os.devnull, 'w'),
-                                     stderr=open(os.devnull, 'w'),
-                                     stdin=open(os.devnull))
-                try:
-                    p.wait()
-                except KeyboardInterrupt:
-                    print()
-                    raise SystemExit
-                try:
-                    hostname = '{} ({})'.format(
-                        ip,
-                        socket.gethostbyaddr(str(ip))[0],
-                    )
-                except socket.herror:
-                    hostname = ip
-                state = 'in use' if p.returncode == 0 else 'not responding'
-                print('{} is {}'.format(hostname, state))
+            sys.stdout.write('{0} is available\n'.format(host))
+
+
+class main(object):
+
+    def _parse_args(self):
+        parser = argparse.ArgumentParser()
+        parser.add_argument('--quiet', '-q', action='store_true',
+                            help='Only print available addresses')
+        parser.add_argument('--num-threads', type=int, default=10,
+                            metavar='COUNT',
+                            help='Number of simultaneous pings to run')
+        parser.add_argument('networks', nargs='+', metavar='network',
+                            help='Network to ping, in CIDR notation')
+        return parser.parse_args()
+
+    def __init__(self):
+        args = self._parse_args()
+        self.queuer = Queuer(h for n in args.networks
+                             for h in ipaddr.IPv4Network(n).iterhosts())
+        self.pingers = []
+        signal.signal(signal.SIGINT, self.quit)
+        self.queuer.start()
+        for __ in range(0, args.num_threads):
+            pinger = Pinger(self.queuer.queue)
+            pinger.quiet = args.quiet
+            self.pingers.append(pinger)
+            pinger.start()
+        while len(threading.enumerate()) > 1:
+            time.sleep(1)
+
+    def quit(self, signum, frame):
+        self.queuer.stop()
+        for pinger in self.pingers:
+            pinger.stop()
 
 
 if __name__ == '__main__':