#!/usr/bin/env python
#
# python nmap wrapper module.  Includes nmap, nmapResult, and nmapException.  
#
# See Docstrings for examples.
#
# TODO:
# 	Add unit testing

import tempfile, os, xml.dom.minidom

class nmap(object):
	"""
	Create and run a nmap scan.  

	Example usage:

		nm = nmap.nmap(interface="eth1", extraOptions="-F -r", sourcePort=20)
		nm.hosts=hostList                       # List of hosts to scan
		nm.ports=[20, 21, 135-139]              # Specify ports in nmap format, e.g. 135-139
		nm.scanType="syn"                       # Implies protocol
		nm.aggressiveness=4
		nm.fingerprint=1                        # Run service and OS fingerprinting
		nm.outputXml="/tmp/nmap-output.xml"     # If not set, will default to a temp file
		nmapResults = nm.scan()                 # Run the scan

	Hosts is the the only required parameter. Absence of any value will
        cause nmap to use its default (except host).  Any value can be passed in during
	object creation, or changed directly with the object instance.
	"""

	# shared static data
	types={ "syn" :     "-sS", 
	        "connect" : "-sT",
		"ack" :     "-sA",
		"window" :  "-sW",
		"maimon" :  "-sM",
		"udp" :     "-sU",
		"null" :    "-sN",
		"fin" :     "-sF",
		"xmas" :    "-sX" }

	# Getters and Setters -- Most of these will be used indirectly by
	# the property() objects, not by module users.
	def getHosts(self):              return self._hosts
	def getPorts(self):              return self._ports
	def getAggressiveness(self):     return self._aggr
	def getExtraOptions(self):       return self._opts
	def getInterface(self):          return self._int
	def getKeepOutput(self):         return self._keepOutput
	def getOutputXml(self):          return self._outputXml
	def getScanType(self):           return self._type
	def getFingerPrint(self):        return self._fp
	def getSourcePort(self):         return self._source

	def getScanCmd(self):
		"""
		Returns a string with the nmap command a scan() will run
		"""
		self._scancmd="nmap "

		# Create two temp files since we need to know where to store/send data
		if not self._inputPath:
			(self._inputFh, self._inputPath) = tempfile.mkstemp(".txt", "nmapHosts_")
			# get file object from FD
			self._inputFh = os.fdopen(self._inputFh, "w")

		# Set output to temporary file
		if not self._outputPath and not self._outputXml:
			(self._outputFh, self._outputPath) = tempfile.mkstemp(".xml", "nmapXml_")
			self._outputFh = os.fdopen(self._outputFh, "w")
		# Set output to the specified output file
		elif not self._outputPath and self._outputXml:
			self._outputPath=self._outputXml

		if self._aggr:
			self._scancmd += "-T %s " % self._aggr
		if self._opts:
			self._scancmd += "%s "    % self._opts
		if self._int:
			self._scancmd += "-i %s " % self._int
		if self._type:
			self._scancmd += "%s "    % nmap.types[self._type]
		if self._fp:
			self._scancmd += "-sV -O "
		if self._source:
			self._scancmd += "-g %s "  % self._source
		if self._ports:
			self._scancmd += "-p"
			for port in self._ports:
				self._scancmd += "%s," % port

			# cut off last char, and add space
			self._scancmd = self._scancmd[:-1] + " "

		if self._hosts:
			self._scancmd += "-iL %s " % self._inputPath
		else:
			raise nmapException("No Hosts defined, at least one required")

		# Send XML output here
		self._scancmd += "-oX %s" % self._outputPath

		# ignore stderr 
		self._scancmd += " >/dev/null 2>&1"

		return self._scancmd

	def setHosts(self, hosts):
		"""
		Set the list of hosts to scan based on a List input
		"""
		self._hosts=hosts

	def delHosts(self, hosts):
		"""
		Remove the hosts in the input List from the scanning queue
		"""
		for host in hosts:
			if host in self._hosts:
				self._hosts.remove(host)

	def setPorts(self, ports):
		"""
		Set the list of ports to scan based on a List input
		"""
		# reset list and assign new values
		self._ports=ports

	def setAggressiveness(self, aggr):
		"""
		How aggresive you want nmap to scan.
		Valid values are >=0 and <=5
		"""
		if aggr and aggr >= 0 and aggr <= 5:
			self._aggr=aggr
		else:
			raise nmapException("Invalid aggressiveness value [%]. Should be >=0 and <=5" % aggr)

		return self._aggr

	def setExtraOptions(self, opts=None):
		"""
		Any nmap command line option you want to add
		"""
		if opts: self._opts=opts
		return self._opts

	def setInterface(self, interface=None):
		if int: self._int=interface
		return self._int

	def setKeepOutput(self, keepOutput=None):
		"""
		Do not delete the output XML after a nmap run
		"""
		if keepOutput: self._keepOutput=keepOutput
		return self._keepOutput

	def setOutputXml(self, outputXml=None):
		if outputXml and not os.path.exists(outputXml):
			self._outputXml=outputXml
		elif outputXml and os.path.exists(outputXml):
			raise nmapException("Output file [%s] already exists" % outputXml)
		return self._outputXml

	def scan(self):
		"""
		Run the nmap scan.  Returns a nmapResults()
		"""
		self._scancmd=self.getScanCmd()

		# Set results in case of multiple scans
		self._results=nmapResult(command=self._scancmd, outputXml=self._outputXml)

		# Write out the hosts file
		[self._inputFh.write(line + "\n") for line in self._hosts]
		self._inputFh.close()
		# Close the output FH so nmap can write to it
		if not self._outputXml:
			self._outputFh.close()

		# run nmap
		rc = os.system(self._scancmd)
		# Get actual return code
		rc = rc >> 8
		if rc == 127:
			raise nmapException("nmap not found in path")
		elif rc >= 1:
			raise nmapException("Execution of nmap failed, rc was [%s]" % rc)

		# parse xml
		doc = xml.dom.minidom.parse(self._outputPath)
		for host in doc.getElementsByTagName("host"):
			address = host.getElementsByTagName("address")[0]
			address = address.getAttribute("addr")

			result = nmapHost(address)

			ports = host.getElementsByTagName("ports")[0]
			ports = ports.getElementsByTagName("port")

			for port in ports:      
				fp = None
				portid=port.getAttribute("portid")
				protocol=port.getAttribute("protocol")

				stateTag = port.getElementsByTagName("state")[0]
				state = stateTag.getAttribute("state")

				# Build up the service and os fingerprint strings
				if self._fp:
					fp=""
					fpTag = port.getElementsByTagName("service")[0]
					if fpTag.hasAttribute("product"):
						fp += fpTag.getAttribute("product") + " "
					if fpTag.hasAttribute("version"):
						fp += fpTag.getAttribute("version") + " "
					if fpTag.hasAttribute("extrainfo"):
						fp += fpTag.getAttribute("extrainfo") 
					if fp == "": fp=None

				# record the port info
				result.ports += [nmapPort(portid, protocol, fp, state, self.sourcePort)]

			# Get OS fingerprint
			osFp = {}
			if self._fp:
				osElement = host.getElementsByTagName("os")[0]
				matchCount=len(osElement.getElementsByTagName("osmatch"))

				# Get all matches
				if matchCount >= 1:
					for match in osElement.getElementsByTagName("osmatch"):
						accuracy = match.getAttribute("accuracy")
						name  = match.getAttribute("name")
						osFp[name]=accuracy

				# Get the best guess if no exact match...
				elif len(osElement.getElementsByTagName("osclass")) > 0:
					for match in osElement.getElementsByTagName("osclass"):
						curFp=""
						accuracy = match.getAttribute("accuracy")
						if match.hasAttribute("vendor"): 
							curFp += match.getAttribute("vendor") + " "
						if match.hasAttribute("osfamily"): 
							curFp += match.getAttribute("osfamily") + " "
						if match.hasAttribute("osgen"): 
							curFp += match.getAttribute("osgen")

						# Add every guess to the os dictionary
						osFp["Guess: %s" % curFp.rstrip()] = accuracy

			result.os=osFp
			self._results.append(result)

		# We should be done with these files now that the scan is finished
		if os.path.exists(self._inputPath): os.unlink(self._inputPath)
		if os.path.exists(self._outputPath) and not self._keepOutput: os.unlink(self._outputPath)
		# Set these so we'll create new ones if needed
		self._inputPath=None
		self._outputPath=None

		return self._results
		
	def setScanType(self, type=None):
		if type in nmap.types:
			self._type=type
		else:
			raise nmapException("Invalid nmap scan type [%s]" % type)

		return self._type

	def setFingerPrint(self, fingerprint=None):
		if fingerprint:
			self._fp=fingerprint
		return self._fp

	def setSourcePort(self, source=None):
		if source:
			self._source=source
		return self._source

	def __repr__(self):
		print self.getScanCmd()

	def __str__(self):
		print self.getScanCmd()

	def __init__(self, hosts=None, ports=None, aggressiveness=None, extraOptions=None, interface=None, 
	             scanType=None, fingerprint=None, sourcePort=None, keepOutput=None, outputXml=None):

		self._results=None
		self._hosts=[]
		self._ports=[]
		self._opts=None
		self._aggr=None
		self._int=None
		self._type=None
		self._fp=None
		self._source=None
		self._inputPath=None
		self._inputFh=None
		self._outputPath=None
		self._outputFh=None
		self._outputXml=None
		self._keepOutput=0
		self._scancmd=None

		if hosts:              self.setHosts(hosts)
		if ports:              self.setPorts(ports)
		if aggressiveness:     self.setAggressiveness(aggressiveness)
		if extraOptions:       self.setExtraOptions(extraOptions)
		if interface:          self.setInterface(interface)
		if keepOutput:         self.setKeepOutput(keepOutput)
		if outputXml:          self.setOutputXml(outputXml)
		if scanType:           self.setScanType(scanType)
		if fingerprint:        self.setFingerPrint(fingerprint)
		if sourcePort:         self.setSourcePort(sourcePort)

	hosts =              property(getHosts, setHosts, delHosts, None)
	ports =              property(getPorts, setPorts, None, None)
	aggressiveness =     property(getAggressiveness, setAggressiveness, None, None)
	extraOptions =       property(getExtraOptions, setExtraOptions, None, None)
	interface =          property(getInterface, setInterface, None, None)
	keepOutput =         property(getKeepOutput, setKeepOutput, None, None)
	outputXml =          property(getOutputXml, setOutputXml, None, None)
	scanType =           property(getScanType, setScanType, None, None)
	fingerprint =        property(getFingerPrint, setFingerPrint, None, None)
	sourcePort =         property(getSourcePort, setSourcePort, None, None)


class nmapHost(object):
	"""
	Part of a nmapResult()
	"""
	def getHostname(self): return self._hostname
	def getOs(self): return self._os
	def getPorts(self, state=None): 
		# return only ports with the state we care about
		if state:
			results=[]
			for port in self._ports:
				if port.state == state: results.append(port)

			return results
		else:
			return self._ports

	def setPorts(self, ports):
		self._ports=ports
		return self._ports

	def setOs(self, os):
		self._os=os
		return self._os

	def setHostname(self, hostname):
		self._hostname=hostname
		return self._hostname

	def __init__(self, hostname):
		self._ports=[]
		self._hostname=hostname
		self._os=None

	hostname = property(getHostname, setHostname, None, None)
	# OS is a dictionary of all possible matches or guesses
	os =       property(getOs, setOs, None, None)
	ports =    property(getPorts, setPorts, None, None)

class nmapPort(object):
	"""
	Part of a nmapHost(), which is part of a nmapResult()
	"""

	# since we don't need to modify the getter/setters, we'll just use __slots__
	__slots__ = ['port', 'fingerprint', 'protocol', 'state', 'sourcePort']

	def __init__(self, port, protocol, fingerprint, state, sourcePort=None):
		self.fingerprint=fingerprint
		self.port=port
		self.state=state
		self.protocol=protocol
		self.sourcePort=sourcePort

class nmapResult(list):
	"""
	Results from a nmap.scan()

	Example usage:
		nmapResults = nmap.scan()

		for host in nmapResults:
		      print "Host: %s " % host.hostname
		      if host.os:
			      for key in host.os.keys():
				      print "OS: %s accuracy %s" % (key, host.os[key])
		      for port in host.getPorts("open"):
			      print "  Open Port: %s %s %s" % (port.protocol, port.port, port.state)
			      if port.fingerprint: print "  Service fingerprint: %s " % port.fingerprint
	"""

	__slots__ = ['command', 'output', 'outputXml']

	def __init__(self, command=None, outputXml=None):
		# This is the command the nmap was called with
		self.command=command
		self.output=""
		self.outputXml=outputXml
		list.__init__(self)

	def __str__(self):
		"""
		Print a basic output report from the results
		"""

		self.output=""
		if self.command:
			self.output = " [*] nmap command was [%s]" % self.command
		for host in self:
			# ugh. report-like output code always looks like crap.
			self.output += "\n [*] Host: [%s]" % host.hostname
			if host.os: 
				for key in host.os.keys():
					self.output += "\n     OS: [%s] accuracy [%s]" % (key, host.os[key]) 

			for port in host.getPorts("open"):
				self.output += "\n        Port: %s %s %s" % (port.port, port.protocol, port.state)
				if port.fingerprint: self.output += "\n          Service fingerprint: %s " % port.fingerprint
		return self.output

	def __repr__(self): 
		return self.__str__()

class nmapException(Exception): 
	"""
	Exceptions within the nmap module
	"""
	pass


if __name__=="__main__":
	"""
	A test of the nmap module against a few ports on 127.0.0.1
	"""

	print " [*] Running Nmap module test"

	nm=nmap(hosts=["127.0.0.1"], ports=[22, 80, "135-139", 443], aggressiveness=4, extraOptions="-n", 
	        interface=None, scanType="syn", fingerprint=1, sourcePort=20)
	
	print " [*] Nmap test command is: [%s]" % nm.getScanCmd()
	print " [*] Running nmap scan"
	nmapResults = nm.scan()

	print nmapResults

	print " [*] Done"
