local dnsbl = require "dnsbl"
local ipOps = require "ipOps"
local stdnse = require "stdnse"
local table = require "table"

description = [[
Checks target IP addresses against multiple DNS anti-spam and open
proxy blacklists and returns a list of services for which an IP has been flagged.  Checks may be limited by service category (eg: SPAM,
PROXY) or to a specific service name.  ]]

---
-- @usage
-- nmap --script dns-blacklist --script-args='dns-blacklist.ip=<ip>'
-- or
-- nmap -sn <ip> --script dns-blacklist
--
-- @output
-- Pre-scan script results:
-- | dns-blacklist:
-- | 1.2.3.4
-- |   PROXY
-- |     dnsbl.tornevall.org - PROXY
-- |       IP marked as "abusive host".
-- |       Proxy is working
-- |       Proxy has been scanned
-- |   SPAM
-- |     dnsbl.inps.de - SPAM
-- |       Spam Received See: http://www.sorbs.net/lookup.shtml?1.2.3.4
-- |     l2.apews.org - SPAM
-- |     list.quorum.to - SPAM
-- |     bl.spamcop.net - SPAM
-- |_    spam.dnsbl.sorbs.net - SPAM
--
-- Supported blacklist list mode (--script-args dns-blacklist.list):
-- | dns-blacklist:
-- |   PROXY
-- |     socks.dnsbl.sorbs.net
-- |     http.dnsbl.sorbs.net
-- |     misc.dnsbl.sorbs.net
-- |     dnsbl.tornevall.org
-- |   SPAM
-- |     dnsbl.inps.de
-- |     bl.nszones.com
-- |     l2.apews.org
-- |     list.quorum.to
-- |     all.spamrats.com
-- |     bl.spamcop.net
-- |     spam.dnsbl.sorbs.net
-- |_    sbl.spamhaus.org
--
-- @args dns-blacklist.ip string containing the IP to check only needed if
--       running the script as a prerule.
--
-- @args dns-blacklist.mode string containing either "short" or "long"
--       long mode can sometimes provide additional information to why an IP
--       has been blacklisted. (default: long)
--
-- @args dns-blacklist.list lists all services that are available for a
--       certain category.
--
-- @args dns-blacklist.services string containing a comma-separated list of
--       services to query. (default: all)
--
-- @args dns-blacklist.category string containing the service category to query
--       eg. spam or proxy (default: all)
--
--

author = "Patrik Karlsson"
license = "Same as Nmap--See https://nmap.org/book/man-legal.html"
categories = {"external", "safe"}


-- The script can be run either as a host- or pre-rule
hostrule = function() return true end
prerule = function() return true end

local arg_IP       = stdnse.get_script_args(SCRIPT_NAME .. ".ip")
local arg_mode     = stdnse.get_script_args(SCRIPT_NAME .. ".mode") or "long"
local arg_list     = stdnse.get_script_args(SCRIPT_NAME .. ".list")
local arg_services  = stdnse.get_script_args(SCRIPT_NAME .. ".services")
local arg_category  = stdnse.get_script_args(SCRIPT_NAME .. ".category") or "all"

local function listServices()
  local result = {}
  if ( "all" == arg_category ) then
    for cat in pairs(dnsbl.SERVICES) do
      local helper = dnsbl.Helper:new(cat, arg_mode)
      local cat_res= helper:listServices()
      cat_res.name = cat
      table.insert(result, cat_res)
    end
  else
    result = dnsbl.Helper:new(arg_category, arg_mode):listServices()
  end
  return stdnse.format_output(true, result)
end

local function formatResult(result)
  local output = {}
  for _, svc in ipairs(result) do
    if ( svc.result.details ) then
      svc.result.details.name = ("%s - %s"):format(svc.name, svc.result.state)
      table.insert(output, svc.result.details)
    else
      table.insert(output, ("%s - %s"):format(svc.name, svc.result.state))
    end
  end
  return output
end

local function fail (err) return stdnse.format_output(false, err) end

dnsblAction = function(host)

  local helper
  if ( arg_services and ( not(arg_category) or "all" == arg_category:lower() ) ) then
    return fail("A service filter can't be used without a specific category")
  elseif( "all" ~= arg_category ) then
    helper = dnsbl.Helper:new(arg_category, arg_mode)
    helper:setFilter(arg_services)
    local status, err = helper:validateFilter()
    if ( not(status) ) then
      return fail(("%s"):format(err))
    end
  end

  local output = {}
  if ( helper ) then
    local result = helper:checkBL(host.ip)
    if ( #result == 0 ) then return end
    output = formatResult(result)
  else
    for cat in pairs(dnsbl.SERVICES) do
      helper = dnsbl.Helper:new(cat, arg_mode)
      local result = helper:checkBL(host.ip)
      local out_part = formatResult(result)
      if ( #out_part > 0 ) then
        out_part.name = cat
        table.insert(output, out_part)
      end
    end
    if ( #output == 0 ) then return end
  end

  if ( "prerule" == SCRIPT_TYPE ) then
    output.name = host.ip
  end

  return stdnse.format_output(true, output)
end


-- execute the action function corresponding to the current rule
action = function(...)

  if ( arg_mode ~= "short" and arg_mode ~= "long" ) then
    return fail("Invalid argument supplied, mode should be either 'short' or 'long'")
  end

  if ( arg_IP and not(ipOps.todword(arg_IP)) ) then
    return fail("Invalid IP address was supplied")
  end

  -- if the list argument was given, just list the services and abort
  if ( arg_list ) then
    return listServices()
  end

  if ( arg_IP and "prerule" == SCRIPT_TYPE ) then
    return dnsblAction( { ip = arg_IP } )
  elseif ( "hostrule" == SCRIPT_TYPE ) then
    return dnsblAction(...)
  end

end