local ipOps = require "ipOps"
local nmap = require "nmap"
local ssh1 = require "ssh1"
local stdnse = require "stdnse"
local string = require "string"
local table = require "table"

description = [[
Attempts to discover multihomed systems by analysing and comparing
information collected by other scripts. The information analyzed
currently includes, SSL certificates, SSH host keys, MAC addresses,
and Netbios server names.

In order for the script to be able to analyze the data it has dependencies to
the following scripts: ssl-cert,ssh-hostkey,nbtstat.

One or more of these scripts have to be run in order to allow the duplicates
script to analyze the data.
]]

---
-- @usage
-- sudo nmap -PN -p445,443 --script duplicates,nbstat,ssl-cert <ips>
--
-- @output
-- | duplicates: 
-- |   ARP
-- |       MAC: 01:23:45:67:89:0a
-- |           192.168.99.10
-- |           192.168.99.11
-- |   Netbios
-- |       Server Name: WIN2KSRV001
-- |           192.168.0.10
-- |_          192.168.1.10
--


--
-- While the script provides basic duplicate functionality, here are some ideas
-- on improvements.
--
-- Possible additional information sources:
-- * Microsoft SQL Server instance names (Match hostname, version, instance
--   names and ports) - Reliable given several instances
-- * Oracle TNS names - Not very reliable
--
-- Possible enhancements:
-- * Compare hosts across information sources and create a global category
--   in which system duplicates are reported based on more than one source.
-- * Add a reliability index for each information source that indicates how
--   reliable the duplicate match was. This could be an index compared to
--   other information sources as well as an indicator of how good the match
--   was for a particular information source.

author = "Patrik Karlsson"
license = "Same as Nmap--See http://nmap.org/book/man-legal.html"
categories = {"safe"}
dependencies = {"ssl-cert", "ssh-hostkey", "nbstat"}


hostrule = function() return true end
postrule = function() return true end

--- check for the presence of a value in a table
--@param tab the table to search into
--@param item the searched value
--@return a boolean indicating whether the value has been found or not
local function contains(tab, item)
	for _, val in pairs(tab) do
		if val == item then
			return true
		end
	end
	return false
end

local function processSSLCerts(tab)
	
	-- Handle SSL-certificates
	-- We create a new table using the SHA1 digest as index
	local ssl_certs = {}
	for host, v in pairs(tab) do
		for port, sha1 in pairs(v) do
			ssl_certs[sha1] = ssl_certs[sha1] or {}
			if ( not contains(ssl_certs[sha1], host.ip) ) then
				table.insert(ssl_certs[sha1], host.ip)
			end
		end
	end
	
	local results = {}
	for sha1, hosts in pairs(ssl_certs) do
		table.sort(hosts, function(a, b) return ipOps.compare_ip(a, "lt", b) end)
		if ( #hosts > 1 ) then
			table.insert(results, { name = ("Certficate (%s)"):format(sha1), hosts } )
		end
	end
	
	return results
end

local function processSSHKeys(tab)

	local hostkeys = {}

	-- create a reverse mapping key_fingerprint -> host(s)
	for ip, keys in pairs(tab) do
		for _, key in ipairs(keys) do
			local fp = ssh1.fingerprint_hex(key.fingerprint, key.algorithm, key.bits)
			if not hostkeys[fp] then
				hostkeys[fp] = {}
			end
			-- discard duplicate IPs
			if not contains(hostkeys[fp], ip) then
				table.insert(hostkeys[fp], ip)
			end
		end
	end

	-- look for hosts using the same hostkey
	local results = {}
	for key, hosts in pairs(hostkeys) do
		if #hostkeys[key] > 1 then
			table.sort(hostkeys[key], function(a, b) return ipOps.compare_ip(a, "lt", b) end)
			local str = 'Key ' .. key .. ':'
			table.insert( results, { name = str, hostkeys[key] } )
		end
	end

	return results
end

local function processNBStat(tab)

	local results, mac_table, name_table = {}, {}, {}
	for host, v in pairs(tab) do
		mac_table[v.mac] = mac_table[v.mac] or {}
		if ( not(contains(mac_table[v.mac], host.ip)) ) then
			table.insert(mac_table[v.mac], host.ip)
		end

		name_table[v.server_name] = name_table[v.server_name] or {}
		if ( not(contains(name_table[v.server_name], host.ip)) ) then
			table.insert(name_table[v.server_name], host.ip)
		end
	end
	
	for mac, hosts in pairs(mac_table) do
		if ( #hosts > 1 ) then
			table.sort(hosts, function(a, b) return ipOps.compare_ip(a, "lt", b) end)
			table.insert(results, { name = ("MAC: %s"):format(mac), hosts })
		end
	end
	
	for srvname, hosts in pairs(name_table) do
		if ( #hosts > 1 ) then
			table.sort(hosts, function(a, b) return ipOps.compare_ip(a, "lt", b) end)
			table.insert(results, { name = ("Server Name: %s"):format(srvname), hosts })
		end
	end
	
	return results
end

local function processMAC(tab)
	
	local function format_mac(mac)
		local octets = {}
		for _, v in ipairs({ string.byte(mac, 1, #mac) }) do
			octets[#octets + 1] = string.format("%02x", v)
		end
		return stdnse.strjoin(":", octets)
	end
	
	local mac
	local mac_table = {}
	
	for host in pairs(tab) do
		if ( host.mac_addr ) then
			mac = format_mac(host.mac_addr)
			mac_table[mac] = mac_table[mac] or {}
			if ( not(contains(mac_table[mac], host.ip)) ) then
				table.insert(mac_table[mac], host.ip)
			end
		end
	end
	
	local results = {}
	for mac, hosts in pairs(mac_table) do
		if ( #hosts > 1 ) then
			table.sort(hosts, function(a, b) return ipOps.compare_ip(a, "lt", b) end)
			table.insert(results, { name = ("MAC: %s"):format(mac), hosts })
		end
	end
	
	return results
end

postaction = function()
	
	local handlers = {
		['ssl-cert'] = { func = processSSLCerts, name = "SSL" },
		['sshhostkey'] = { func = processSSHKeys, name = "SSH" },
		['nbstat'] = { func = processNBStat, name = "Netbios" },
		['mac'] = { func = processMAC, name = "ARP" }
	}
	
	-- temporary re-allocation code for SSH keys
	for k, v in pairs(nmap.registry.sshhostkey or {}) do
		nmap.registry['duplicates'] = nmap.registry['duplicates'] or {}
		nmap.registry['duplicates']['sshhostkey'] = nmap.registry['duplicates']['sshhostkey'] or {}
		nmap.registry['duplicates']['sshhostkey'][k] = v
	end
	
	if ( not(nmap.registry['duplicates']) ) then
		return
	end
	
	local results = {}
	for key, handler in pairs(handlers) do
		if ( nmap.registry['duplicates'][key] ) then
			local result_part = handler.func( nmap.registry['duplicates'][key] )
			if ( result_part and #result_part > 0 ) then
				table.insert(results, { name = handler.name, result_part } )
			end
		end
	end
	
	return stdnse.format_output(true, results)
end

-- we have no real action in here. In essence we move information from the
-- host based registry to the global one, so that our postrule has access to
-- it when we need it.
hostaction = function(host)

	nmap.registry['duplicates'] = nmap.registry['duplicates'] or {}

	for port, cert in pairs(host.registry["ssl-cert"] or {}) do
		nmap.registry['duplicates']['ssl-cert'] = nmap.registry['duplicates']['ssl-cert'] or {}
		nmap.registry['duplicates']['ssl-cert'][host] = nmap.registry['duplicates']['ssl-cert'][host] or {}
		nmap.registry['duplicates']['ssl-cert'][host][port] = stdnse.tohex(cert:digest("sha1"), { separator = " ", group = 4 })
	end
	
	if ( host.registry['nbstat'] ) then
		nmap.registry['duplicates']['nbstat'] = nmap.registry['duplicates']['nbstat'] or {}
		nmap.registry['duplicates']['nbstat'][host] = host.registry['nbstat']
	end

	if ( host.mac_addr_src ) then
		nmap.registry['duplicates']['mac'] = nmap.registry['duplicates']['mac'] or {}
		nmap.registry['duplicates']['mac'][host] = true
	end

	return
end

local Actions = {
  hostrule = hostaction,
  postrule = postaction
}

-- execute the action function corresponding to the current rule
action = function(...) return Actions[SCRIPT_TYPE](...) end