Below is the file 'scan.py' from this revision. You can also download the file.
#!/usr/bin/python from goatpy.nmapwrapper import nmap from goatpy.gen2consume import gen2consume import ConfigParser import threading import select import popen2 import fcntl import time import sets import sys import os import config # urgency urgency_info = 0 urgency_notice = 1 urgency_warning = 2 urgency_alert = 3 urgency_critical = 4 urgency_to_string = { urgency_info : 'Info', urgency_notice : 'Notice', urgency_warning : 'Warning', urgency_alert : 'Alert', urgency_critical : 'CRITICAL' } # plugins should provide a host_callback method that takes args: # scanner : an instance of Scanner (as below) # host : Host() instance # tcp_ports : open, interesting ports for this plugin # udp_ports : open, interesting ports for this plugin # they should return a list of ScannerResponse objects def set_nonblocking(fd): fl = fcntl.fcntl(fd, fcntl.F_GETFL) fcntl.fcntl(fd, fcntl.F_SETFL, fl | os.O_NDELAY) def scan_host(mythread, host, library): library.notify("[%2d] %-15s: has been assigned (T:%s U:%s)" % (mythread.id, host.addresses[0][0], str(host.tcp_ports), str(host.udp_ports))) plugins = library.plugins_for_ports(host.tcp_ports, host.udp_ports) errors = [] results = [] # if the magic UDP port is open, then ignore any open UDP ports if config.magic_udp_port in host.udp_ports: library.notify("ignoring UDP ports; magic UDP port %d open" % (config.magic_udp_port)) host.udp_ports = [] for plugin in plugins: tcp_ports = filter(lambda x: x in plugins[plugin][0], host.tcp_ports) udp_ports = filter(lambda x: x in plugins[plugin][1], host.udp_ports) try: library.notify("[%2d] %-15s: run plugin %s" % (mythread.id, host.addresses[0][0], plugin.export['name'])) mythread.status_string = "in plugin %s" % (plugin.export['name']) result = plugin.host_callback(library, host, tcp_ports, udp_ports) mythread.status_string = "completed plugin %s" % (plugin.export['name']) results += result except: import traceback t,v,tr = sys.exc_info() err_data = '\n'.join(traceback.format_exception(t, v, tr)) errors.append((plugin.export['name'], err_data)) mythread.status_string = "writing results" for plugin in library.output_plugins: plugin.write_results(library, host, results) plugin.write_errors(library, host, errors) mythread.status_string = None library.notify("[%2d] %-15s: done with this host" % (mythread.id, host.addresses[0][0])) # (I hate Unix) : setting a SIGCHLD handler will start me getting # lots of exceptions whenver I do blah.read(), select(), etc.. which # makes the code of this program so _massively_ ugly that I'll just # call wait occassionally and it'll probably even out # # You can't just call wait() in a thread that does a SIGCHLD because # Unix won't let you. Damnit. class WaitThread(threading.Thread): def __init__(self): for base in self.__class__.__bases__: base.__init__(self) self.setDaemon(True) def run(self): while 1: try: procs = os.wait() print "Processes have terminated:", procs except: pass import time time.sleep(1) class PluginLibrary: def __init__(self, config): self.config = config orig_path = None self.output_plugins = [] self.tcp_to_plugin = {} self.udp_to_plugin = {} orig_path = sys.path sys.path.insert(0, self.config.plugin_path) for plugin in self.config.plugins: mod = self.do_load_plugin(plugin) if mod.export['type'] == "scanner": self.link_ports(mod, mod.export['tcp_ports'], self.tcp_to_plugin) self.link_ports(mod, mod.export['udp_ports'], self.udp_to_plugin) elif mod.export['type'] == "output": mod.initialise(self) self.output_plugins.append(mod) if orig_path: sys.path = orig_path def finalise_plugins(self): map(lambda x: x.finalise(self), self.output_plugins) def notify(self, str): print "%s" % (str) def link_ports(self, mod, port_list, hash): if not port_list: return for port in port_list: if not hash.has_key(port): hash[port] = [] hash[port].append(mod) def get_ports(self): return self.tcp_to_plugin.keys(), self.udp_to_plugin.keys() def do_load_plugin(self, plugin_name): mod = __import__('%s' % plugin_name, globals(), locals(), ['']) # self.notify("Plugin loaded: " + mod.export['name'] + " version " + mod.export['version'] + " (" + mod.export['type'] + ")") return mod def plugins_for_ports(self, tcp_ports, udp_ports): rv = {} for port in tcp_ports: for module in self.tcp_to_plugin.get(port, []): if not rv.has_key(module): rv[module] = ([], []) rv[module][0].append(port) for port in udp_ports: for module in self.udp_to_plugin.get(port, []): if not rv.has_key(module): rv[module] = ([], []) rv[module][1].append(port) return rv class ScannerResponse: def __init__(self): self.address = None self.plugin_name = None self.tcp_ports = None self.udp_ports = None self.urgency = None self.short_mesg = None self.long_mesg = None self.advice = None def urgency_string(self): return urgency_to_string[self.urgency] def text_summary(self): if self.urgency == None or self.plugin_name == None: return ports = [] if self.tcp_ports and len(self.tcp_ports): ports.append("T:" + ','.join([str(t) for t in self.tcp_ports])) if self.udp_ports and len(self.udp_ports): ports.append("U:" + ','.join([str(t) for t in self.udp_ports])) return "%-10s %-10s %-10s %s" % (','.join(ports), self.plugin_name, urgency_to_string[self.urgency], self.short_mesg) class Scanner: def __init__(self, config_path): self._config_path = config_path self.config = config self.library = PluginLibrary(self.config) def notify(self, str): print "%s" % (str) def scan(self): nmap_command = self.config.nmap_command + ' ' + self.config.nmap_options ports = [] tcp_ports, udp_ports = self.library.get_ports() if len(tcp_ports): ports.append("T:" + ','.join([str(t) for t in tcp_ports])) if len(udp_ports): udp_ports.append(config.magic_udp_port) ports.append("U:" + ','.join([str(t) for t in udp_ports])) if not len(ports): self.notify("Nothing to do; no ports requested by plugins.") return nmap_command += ' -p ' + ','.join(ports) nmap_command += ' -oX -' range = ' '.join(self.config.scan_networks) nmap_command += ' ' + range self.notify("Starting scan: %s" % (nmap_command)) gt = gen2consume(nmap(nmap_command), self.config.max_threads, lambda thread, host: scan_host(thread, host, self.library)) gt.join() countdown = self.config.thread_wait_time while countdown > 0: running = filter(lambda x: x.isAlive(), threading.enumerate()) has_status = filter(lambda x: "status_string" in dir(x) and x.status_string != None, running) cnt = len(running) reasons = map(lambda x: x.status_string, has_status) print "[%2d secs] Waiting for %d threads (%s)" % (countdown, cnt, ','.join(reasons)) if cnt == 2: break countdown -= 1 time.sleep(1) print "Finalising output plugins." self.library.finalise_plugins() print "Done!" if __name__ == '__main__': wt = WaitThread() wt.start() scanner = Scanner('scanner.cfg') scanner.scan() sys.exit(0)