local datafiles = require "datafiles"
local netbios = require "netbios"
local nmap = require "nmap"
local stdnse = require "stdnse"
local string = require "string"
local table = require "table"

description = [[
Attempts to retrieve the target's NetBIOS names and MAC address.

By default, the script displays the name of the computer and the logged-in
user; if the verbosity is turned up, it displays all names the system thinks it
owns.
]]

---
-- @usage
-- sudo nmap -sU --script nbstat.nse -p137 <host>
--
-- @output
-- Host script results:
-- | nbstat: 
-- | name: WINDOWS2003
-- | user: <unknown>
-- |_MAC: 00:0c:29:c6:da:f5 (VMware)
--
-- Host script results:
-- | nbstat: 
-- | name: WINDOWS2003
-- | user: <unknown>
-- | MAC: 00:0c:29:c6:da:f5 (VMware)
-- | names: 
-- |   WINDOWS2003<00>      Flags: <unique><active>
-- |   WINDOWS2003<20>      Flags: <unique><active>
-- |   SKULLSECURITY<00>    Flags: <group><active>
-- |   SKULLSECURITY<1e>    Flags: <group><active>
-- |   SKULLSECURITY<1d>    Flags: <unique><active>
-- |_  \x01\x02__MSBROWSE__\x02<01>  Flags: <group><active>
--
--
--@xmloutput
-- <script id="nbstat" output="...">
--   <table key="names">
--     <table>
--       <elem key="suffix">1</elem>
--       <elem key="flags">33792</elem>
--       <elem key="name">\x01\x02__MSBROWSE__\x02</elem>
--     </table>
--     <table>
--       <elem key="suffix">0</elem>
--       <elem key="flags">1024</elem>
--       <elem key="name">WINDOWS2003</elem>
--     </table>
--     <table>
--       <elem key="suffix">3</elem>
--       <elem key="flags">1024</elem>
--       <elem key="name">WINDOWS2003</elem>
--     </table>
--     <table>
--       <elem key="suffix">32</elem>
--       <elem key="flags">1024</elem>
--       <elem key="name">WINDOWS2003</elem>
--     </table>
--     <table>
--       <elem key="suffix">29</elem>
--       <elem key="flags">1024</elem>
--       <elem key="name">SKULLSECURITY</elem>
--     </table>
--     <table>
--       <elem key="suffix">30</elem>
--       <elem key="flags">33792</elem>
--       <elem key="name">SKULLSECURITY</elem>
--     </table>
--     <table>
--       <elem key="suffix">0</elem>
--       <elem key="flags">33792</elem>
--       <elem key="name">SKULLSECURITY</elem>
--     </table>
--   </table>
--   <table key="statistics">
--     <elem>00 0c 29 c6 da f5 00 00 00 00 00 00 00 00 00 00</elem>
--     <elem>00 00 00 00 00 00 00 00 00 00 00 00 00 00 00 00</elem>
--     <elem>00 00 00 00 00 00 00 00 00 00 00 00 00 00</elem>
--   </table>
--   <elem key="name">WINDOWS2003</elem>
--   <table key="MAC">
--     <elem key="address">00:0c:29:c6:da:f5</elem>
--     <elem key="manufacturer">VMware</elem>
--   </table>
--   <elem key="user">&lt;unknown&gt;</elem>
-- </script>


author = "Brandon Enright, Ron Bowes"
license = "Same as Nmap--See http://nmap.org/book/man-legal.html"

-- Current version of this script was based entirly on Implementing CIFS, by 
-- Christopher R. Hertel. 
categories = {"default", "discovery", "safe"}


hostrule = function(host)

	-- The following is an attempt to only run this script against hosts
	-- that will probably respond to a UDP 137 probe.  One might argue
	-- that sending a single UDP packet and waiting for a response is no
	-- big deal and that it should be done for every host.  In that case
	-- simply change this rule to always return true.

	local port_t135 = nmap.get_port_state(host,
		{number=135, protocol="tcp"})
	local port_t139 = nmap.get_port_state(host,
		{number=139, protocol="tcp"})
	local port_t445 = nmap.get_port_state(host,
		{number=445, protocol="tcp"})
	local port_u137 = nmap.get_port_state(host,
		{number=137, protocol="udp"})

	return (port_t135 ~= nil and port_t135.state == "open") or
		(port_t139 ~= nil and port_t139.state == "open") or
		(port_t445 ~= nil and port_t445.state == "open") or
		(port_u137 ~= nil and
			(port_u137.state == "open" or
			port_u137.state == "open|filtered"))
end


action = function(host)

	local i
	local status
	local names, statistics
	local server_name, user_name
	local mac = {}
  local prefix
	local response = {}
	local catch = function() return end
	local try = nmap.new_try(catch)
	

	-- Get the list of NetBIOS names
	status, names, statistics = netbios.do_nbstat(host.ip)
	status, names, statistics = netbios.do_nbstat(host.ip)
	status, names, statistics = netbios.do_nbstat(host.ip)
	status, names, statistics = netbios.do_nbstat(host.ip)
	if(status == false) then
		return stdnse.format_output(false, names)
	end

	-- Get the server name
	status, server_name = netbios.get_server_name(host.ip, names)
	if(status == false) then
		return stdnse.format_output(false, server_name)
	end

	-- Get the logged in user
	status, user_name = netbios.get_user_name(host.ip, names)
	if(status == false) then
		return stdnse.format_output(false, user_name)
	end

  local mac_prefixes = try(datafiles.parse_mac_prefixes())
	
	-- Format the Mac address in the standard way
	if(#statistics >= 6) then
		-- MAC prefixes are matched on the first three bytes, all uppercase
		prefix = string.upper(string.format("%02x%02x%02x", statistics:byte(1), statistics:byte(2), statistics:byte(3)))
		mac.manufacturer = mac_prefixes[prefix]
		mac.address = string.format("%02x:%02x:%02x:%02x:%02x:%02x",
        statistics:byte(1), statistics:byte(2), statistics:byte(3),
        statistics:byte(4), statistics:byte(5), statistics:byte(6)
        )
		host.registry['nbstat'] = { 
			server_name = server_name, 
			mac = mac.address
		}
		-- Samba doesn't set the Mac address, and nmap-mac-prefixes shows that as Xerox
		if(mac.address == "00:00:00:00:00:00") then
			mac.address = "<unknown>"
      mac.manufacturer = nil
		end
	else
		mac.address = "<unknown>"
	end

	-- Check if we actually got a username
	if(user_name == nil) then
		user_name = "<unknown>"
	end


  stdnse.set_tostring(mac, function(m)
      return string.format("%s (%s)", m.address, m.manufacturer or "unknown")
    end)
  local response = { name= server_name, user= user_name, MAC= mac }
  local key_order = {"name", "user", "MAC"}
	-- If verbosity is set, dump the whole list of names
	if(nmap.verbosity() >= 1) then
    local function namestring(name)
			local flags_str = netbios.flags_to_string(name['flags'])
      local name_str = string.format("%s<%02x>", name.name, name.suffix)
      return string.format("%-20s Flags: %s", name_str, flags_str)
    end

    stdnse.set_tostring(names, function(names)
      local out = {''}
      for _,n in ipairs(names) do
        out[#out+1] = namestring(n)
      end
      return table.concat(out, "\n  ")
    end)

    response["names"] = names
    key_order[#key_order+1] = "names"

		-- If super verbosity is set, print out the full statistics
		if(nmap.verbosity() >= 2) then
      local stats = {}
      for i= 1, #statistics, 16 do
        stats[#stats+1] = stdnse.tohex(statistics:sub(i,i+15), {separator=" "})
      end
      response.statistics = stats
      key_order[#key_order+1] = "statistics"
		end

	end

  stdnse.set_tostring(response,
      stdnse.format_generator({
          key_order=key_order
    }))
  return response

end