Python - Can't kill main thread with KeyboardInterrupt

RedSparrow picture RedSparrow · Aug 1, 2013 · Viewed 18.7k times · Source

I'm making a simple multi-threaded port scanner. It scans all ports on host and returns open ports. The trouble is interrupting the scan. It take a lot of time for a scan to complete and sometimes I wish to kill program with C-c while in the middle of scan. Trouble is the scan won't stop. Main thread is locked on queue.join() and oblivious to KeyboardInterrupt, until all data from queue is processed thus deblocking main thread and exiting program gracefully. All my threads are daemonized so when main thread dies they should die with him.

I tried using signal lib, no success. Overriding threading.Thread class and adding method for graceful termination didn't work... Main thread just won't receive KeyboardInterrupt while executing queue.join()

import threading, sys, Queue, socket

queue = Queue.Queue()

def scan(host):
    while True:
        port = queue.get()

        if port > 999 and port % 1000 == 0:
            print port
        try:
            #sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 
            #sock.settimeout(2) #you need timeout or else it will try to connect forever! 
            #sock.connect((host, port))
            #----OR----
            sock = socket.create_connection((host, port), timeout = 2)

            sock.send('aaa')
            data = sock.recv(100)
            print "Port {} open, message: {}".format(port, data)
            sock.shutdown()
            sock.close()
            queue.task_done()
        except:
            queue.task_done()


def main(host):
    #populate queue
    for i in range(1, 65536):
        queue.put(i)
    #spawn worker threads
    for port in range(100):
        t = threading.Thread(target = scan, args = (host,))
        t.daemon = True
        t.start()

if __name__ == '__main__':
    host = ""

    #does input exist?
    try:
        host = sys.argv[1]
    except:
        print "No argument was recivied!"
        exit(1)

    #is input sane?
    try:
        host = socket.gethostbyname(host)
    except:
        print "Adress does not exist"
        exit(2)

    #execute main program and wait for scan to complete
    main(host)
    print "Post main() call!"
    try:
        queue.join()
    except KeyboardInterrupt:
        print "C-C"
        exit(3)

EDIT:

I have found a solution by using time module.

#execute main program and wait for scan to complete
main(host)

#a little trick. queue.join() makes main thread immune to keyboardinterrupt. So use queue.empty() with time.sleep()
#queue.empty() is "unreliable" so it may return True a bit earlier then intented.
#when queue is true, queue.join() is executed, to confirm that all data was processed.
#not a true solution, you can't interrupt main thread near the end of scan (when queue.empty() returns True)
try:
    while True:
        if queue.empty() == False:
            time.sleep(1)
        else:
            break
except KeyboardInterrupt:
    print "Alas poor port scanner..."
    exit(1)
queue.join()

Answer

solusipse picture solusipse · Aug 1, 2013

You made your threads daemons already, but you need to keep your main thread alive while daemon threads are there, there's how to do that: Cannot kill Python script with Ctrl-C