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, library.scanner_names, 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
		self.urgency_to_string = urgency_to_string
		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)
		self.scanner_names = []
		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)
				self.scanner_names.append(mod.export['name'])
			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 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)