require File.join(File.expand_path(File.dirname(__FILE__)),"nokogiri_doc_mixin")

module Rex
	module Parser

		# If Nokogiri is available, define Template document class. 
		load_nokogiri && class FoundstoneDocument < Nokogiri::XML::SAX::Document

		include NokogiriDocMixin

		def start_document
			@report_type_ok = true # Optimistic
		end

		# Triggered every time a new element is encountered. We keep state
		# ourselves with the @state variable, turning things on when we
		# get here (and turning things off when we exit in end_element()).
		def start_element(name=nil,attrs=[])
			attrs = normalize_attrs(attrs)
			block = @block
			return unless @report_type_ok
			@state[:current_tag][name] = true
			case name
			when "ReportInfo"
				check_for_correct_report_type(attrs,&block)
			when "Host"
				record_host(attrs)
			when "Service" 
				record_service(attrs)
			when "Port", "Protocol", "Banner"
				@state[:has_text] = true
			when "Vuln" # under VulnsFound, ignore risk 0 things
				record_vuln(attrs)
			when "Risk" # for Vuln
				@state[:has_text] = true
			when "CVE" # Under Vuln
				@state[:has_text] = true
			end
		end

		# When we exit a tag, this is triggered.
		def end_element(name=nil)
			block = @block
			return unless @report_type_ok
			case name
			when "Host" # Wrap it up
				collect_host_data
				host_object = report_host &block
				if host_object
					db.report_import_note(@args[:wspace],host_object)
					report_fingerprint(host_object)
					report_services(host_object)
					report_vulns(host_object)
				end
				# Reset the state once we close a host
				@state.delete_if {|k| k != :current_tag}
			when "Port" 
				@state[:has_text] = false
				collect_port
			when "Protocol"
				@state[:has_text] = false
				collect_protocol
			when "Banner"
				collect_banner
				@state[:has_text] = false
			when "Service"
				collect_service
			when "Vuln"
				collect_vuln
			when "Risk"
				@state[:has_text] = false
				collect_risk
			when "CVE"
				@state[:has_text] = false
				collect_cve
			end
			@state[:current_tag].delete name
		end

		# Nothing technically stopping us from parsing this as well, 
		# but saving this for later
		def check_for_correct_report_type(attrs,&block)
			report_type = attr_hash(attrs)["ReportType"] 
			if report_type == "Network Inventory"
				@report_type_ok = true
			else
				if report_type == "Risk Data"
					msg = "The Foundstone/Mcafee report type '#{report_type}' is not currently supported"
					msg << ",\nso no data will be imported. Please use the 'Network Inventory' report instead."
				else
					msg = ".\nThe Foundstone/Macafee report type '#{report_type}' is unsupported."
				end
				db.emit(:warning,msg,&block) if block
				@report_type_ok = false
			end
		end

		def collect_risk
			return unless in_tag("VulnsFound")
			return unless in_tag("HostData")
			return unless in_tag("Host")
			risk = @text.to_s.to_i
			@state[:vuln][:risk] = risk
			@text = nil
		end

		def collect_cve
			return unless in_tag("VulnsFound")
			return unless in_tag("HostData")
			return unless in_tag("Host")
			cve = @text.to_s
			@state[:vuln][:cves] ||= []
			@state[:vuln][:cves] << cve unless cve == "CVE-MAP-NOMATCH"
			@text = nil
		end

		# Determines if we should keep the vuln or not. Note that
		# we cannot tie them to a service.
		def collect_vuln
			return unless in_tag("VulnsFound")
			return unless in_tag("HostData")
			return unless in_tag("Host")
			return if @state[:vuln][:risk] == 0
			@report_data[:vulns] ||= []
			vuln_hash = {}
			vuln_hash[:name] = @state[:vuln]["VulnName"]
			refs = []
			refs << "FID-#{@state[:vuln]["id"]}"
			if @state[:vuln][:cves]
				@state[:vuln][:cves].each {|cve| refs << cve}
			end
			vuln_hash[:refs] = refs
			@report_data[:vulns] << vuln_hash
		end

		# These are per host.
		def record_vuln(attrs)
			return unless in_tag("VulnsFound")
			return unless in_tag("HostData")
			return unless in_tag("Host")
			@state[:vulns] ||= []
			
			@state[:vuln] = attr_hash(attrs) # id and VulnName
		end

		def record_service(attrs)
			return unless in_tag("ServicesFound")
			return unless in_tag("Host")
			@state[:service] = attr_hash(attrs)
		end

		def collect_port
			return unless in_tag("Service")
			return unless in_tag("ServicesFound")
			return unless in_tag("Host")
			return if @text.nil? || @text.empty?
			@state[:service][:port] = @text.strip
			@text = nil
		end

		def collect_protocol
			return unless in_tag("Service")
			return unless in_tag("ServicesFound")
			return unless in_tag("Host")
			return if @text.nil? || @text.empty?
			@state[:service][:proto] = @text.strip
			@text = nil
		end

		def collect_banner
			return unless in_tag("Service")
			return unless in_tag("ServicesFound")
			return unless in_tag("Host")
			return if @text.nil? || @text.empty?
			banner = normalize_foundstone_banner(@state[:service]["ServiceName"],@text)
			unless banner.nil? || banner.empty?
				@state[:service][:banner] = banner
			end
			@text = nil
		end

		def collect_service
			return unless in_tag("ServicesFound")
			return unless in_tag("Host")
			return unless @state[:service][:port]
			@report_data[:ports] ||= [] 
			port_hash = {}
		 	port_hash[:port] = @state[:service][:port]
		 	port_hash[:proto] = @state[:service][:proto]
		 	port_hash[:info] = @state[:service][:banner]
		 	port_hash[:name] = db.nmap_msf_service_map(@state[:service]["ServiceName"])
			@report_data[:ports] << port_hash
		end

		def record_host(attrs)
			return unless in_tag("HostData")
			@state[:host] = attr_hash(attrs)
		end

		def collect_host_data
			@report_data[:host] = @state[:host]["IPAddress"]
			if @state[:host]["NBName"] && !@state[:host]["NBName"].empty?
				@report_data[:name] = @state[:host]["NBName"]
			elsif @state[:host]["DNSName"] && !@state[:host]["DNSName"].empty?
				@report_data[:name] = @state[:host]["DNSName"]
			end
			if @state[:host]["OSName"] && !@state[:host]["OSName"].empty?
				@report_data[:os_fingerprint] = @state[:host]["OSName"]
			end
			@report_data[:state] = Msf::HostState::Alive
			@report_data[:mac] = @state[:mac] if @state[:mac]
		end

		def report_host(&block)
			return unless in_tag("HostData")
			if host_is_okay
				db.emit(:address,@report_data[:host],&block) if block
				host_info = @report_data.merge(:workspace => @args[:wspace])
				db_report(:host,host_info)
			end
		end

		def report_fingerprint(host_object)
			fp_note = {
				:workspace => host_object.workspace,
				:host => host_object,
				:type => 'host.os.foundstone_fingerprint',
				:data => {:os => @report_data[:os_fingerprint] }
			}
			db_report(:note, fp_note)
		end

		def report_services(host_object)
			return unless in_tag("HostData")
			return unless host_object.kind_of? Msf::DBManager::Host
			return unless @report_data[:ports]
			return if @report_data[:ports].empty?
			@report_data[:ports].each do |svc|
				db_report(:service, svc.merge(:host => host_object))
			end
		end

		def report_vulns(host_object)
			return unless in_tag("HostData")
			return unless host_object.kind_of? Msf::DBManager::Host
			return unless @report_data[:vulns]
			return if @report_data[:vulns].empty?
			@report_data[:vulns].each do |vuln|
				db_report(:vuln, vuln.merge(:host => host_object))
			end
		end

		# Foundstone's banners are pretty free-form
		# and often not just banners. Clean them up
		# for the :info field, delegate off for other
		# protocol data we can use.
		def normalize_foundstone_banner(service,banner)
			return "" if(banner.nil? || banner.strip.empty?)
			if first_line_only? service
				return (first_line banner)
			elsif needs_more_processing? service
				return process_service(service,banner)
			else
				return (first_line banner)
			end
		end

		# Services where we only care about the first
		# line of the banner tag.
		def first_line_only?(service)
			svcs = %w{ 
				vnc ftp ftps smtp oracle-tns nntp ssh ntp
			}
			9.times {|i| svcs << "vnc-#{i}"}
			svcs.include? service
		end

		# Services where we need to do more processing
		# before handing the banner back. 
		def needs_more_processing?(service)
			svcs = %w{ 
				microsoft-ds loc-srv http https sunrpc netbios-ns
		 	}
			svcs.include? service
		end

		def first_line(str)
			str.split("\n").first.to_s.strip
		end

		# XXX: Actually implement more of these 
		def process_service(service,banner)
			meth = "process_service_#{service.gsub("-","_")}"
			if self.respond_to? meth
				self.send meth, banner
			else
				return (first_line banner)
			end
		end

		# XXX: Register a proper netbios note as the regular
		# scanner does.
		def process_service_netbios_ns(banner)
			mac_regex = /[0-9A-Fa-f:]{17}/
			@state[:mac] = banner[mac_regex]
			first_line banner
		end

		# XXX: Make this behave more like the smb scanner
		def process_service_microsoft_ds(banner)
			lm_regex = /Native LAN Manager/
			lm_banner = nil
			banner.each_line { |line| 
				if line[lm_regex]
					lm_banner = line 
					break
				end
			}
			lm_banner || first_line(banner)
		end

		def process_service_http(banner)
			server = nil
			banner.each_line do |line|
				if line =~ /^Server:\s+(.*)/
					server = $1
					break
				end
			end
			server || first_line(banner)
		end

		alias :process_service_https :process_service_http
		alias :process_service_rtsp :process_service_http

	end

end
end