local bin = require "bin"
local nmap = require "nmap"
local shortport = require "shortport"
local stdnse = require "stdnse"
local string = require "string"
local tab = require "tab"
local table = require "table"

description = [[
Queries Quake3-style master servers for game servers (many games other than Quake 3 use this same protocol).
]]

---
-- @usage
-- nmap -sU -p 27950 --script=quake3-master-getservers <target>
-- @output
-- PORT      STATE SERVICE REASON
-- 27950/udp open  quake3-master
-- | quake3-master-getservers:
-- |     192.0.2.22:26002 Xonotic (Xonotic 3)
-- |     203.0.113.37:26000 Nexuiz (Nexuiz 3)
-- |_    Only 2 shown. Use --script-args quake3-master-getservers.outputlimit=-1 to see all.
--
-- @args quake3-master-getservers.outputlimit If set, limits the amount of
--       hosts returned by the script. All discovered hosts are still
--       stored in the registry for other scripts to use. If set to 0 or
--       less, all files are shown. The default value is 10.

author = "Toni Ruottu"
license = "Same as Nmap--See https://nmap.org/book/man-legal.html"
categories = {"default", "discovery", "safe"}

portrule = shortport.port_or_service ({20110, 20510, 27950, 30710}, "quake3-master", {"udp"})
postrule = function()
  return (nmap.registry.q3m_servers ~= nil)
end

-- There are various sources for this information. These include:
-- - http://svn.icculus.org/twilight/trunk/dpmaster/readme.txt?view=markup
-- - http://openarena.wikia.com/wiki/Changes
-- - http://dpmaster.deathmask.net/
-- - qstat-2.11, qstat.cfg
-- - scanning master servers
-- - looking at game traffic with Wireshark
local KNOWN_PROTOCOLS = {
  ["5"] = "Call of Duty",
  ["10"] = "unknown",
  ["43"] = "unknown",
  ["48"] = "unknown",
  ["50"] = "Return to Castle Wolfenstein",
  ["57"] = "unknown",
  ["59"] = "Return to Castle Wolfenstein",
  ["60"] = "Return to Castle Wolfenstein",
  ["66"] = "Quake III Arena",
  ["67"] = "Quake III Arena",
  ["68"] = "Quake III Arena, or Urban Terror",
  ["69"] = "OpenArena, or Tremulous",
  ["70"] = "unknown",
  ["71"] = "OpenArena",
  ["72"] = "Wolfenstein: Enemy Territory",
  ["80"] = "Wolfenstein: Enemy Territory",
  ["83"] = "Wolfenstein: Enemy Territory",
  ["84"] = "Wolfenstein: Enemy Territory",
  ["2003"] = "Soldier of Fortune II: Double Helix",
  ["2004"] = "Soldier of Fortune II: Double Helix",
  ["DarkPlaces-Quake 3"] = "DarkPlaces Quake",
  ["Nexuiz 3"] = "Nexuiz",
  ["Transfusion 3"] = "Transfusion",
  ["Warsow 8"] = "Warsow",
  ["Xonotic 3"] = "Xonotic",
}

local function getservers(host, port, q3protocol)
  local socket = nmap.new_socket()
  socket:set_timeout(10000)
  local status, err = socket:connect(host, port)
  if not status then
    return {}
  end
  local probe = bin.pack("CCCCA", 0xff, 0xff, 0xff, 0xff, string.format("getservers %s empty full\n", q3protocol))
  socket:send(probe)

  local data
  status, data = socket:receive() -- get some data
  if not status then
    return {}
  end
  nmap.set_port_state(host, port, "open")

  local magic = bin.pack("CCCCA", 0xff, 0xff, 0xff, 0xff, "getserversResponse")
  local tmp
  while #data < #magic do -- get header
    status, tmp = socket:receive()
    if status then
      data = data .. tmp
    end
  end
  if string.sub(data, 1, #magic) ~= magic then -- no match
    return {}
  end

  port.version.name = "quake3-master"
  nmap.set_port_version(host, port)

  local EOT = bin.pack("ACCC", "EOT", 0, 0, 0)
  local pieces = stdnse.strsplit("\\", data)
  while pieces[#pieces] ~= EOT do -- get all data
    status, tmp = socket:receive()
    if status then
      data = data .. tmp
      pieces = stdnse.strsplit("\\", data)
    end
  end

  table.remove(pieces, 1)       --remove magic
  table.remove(pieces, #pieces) --remove EOT

  local servers = {}
  for _, value in ipairs(pieces) do
    local parts = {bin.unpack("CCCC>S", value)}
    if #parts > 5 then
      local o1 = parts[2]
      local o2 = parts[3]
      local o3 = parts[4]
      local o4 = parts[5]
      local p = parts[6]
      table.insert(servers, {string.format("%d.%d.%d.%d", o1, o2, o3, o4), p})
    end
  end
  socket:close()
  return servers
end

local function formatresult(servers, outputlimit, protocols)
  local t = tab.new()

  if not outputlimit then
    outputlimit = #servers
  end
  for i = 1, outputlimit do
    if not servers[i] then
      break
    end
    local node = servers[i]
    local protocol = node.protocol
    local ip = node.ip
    local portnum = node.port
    tab.addrow(t, string.format('%s:%d', ip, portnum), string.format('%s (%s)', protocols[protocol], protocol))
  end

  return tab.dump(t)
end

local function dropdupes(tables, stringify)
  local unique = {}
  local dupe = {}
  local s
  for _, v in ipairs(tables) do
    s = stringify(v)
    if not dupe[s] then
      table.insert(unique, v)
      dupe[s] = true
    end
  end
  return unique
end

local function scan(host, port, protocols)
  local discovered = {}
  for protocol, _ in pairs(protocols) do
    for _, node in ipairs(getservers(host, port, protocol)) do
      local entry = {
        protocol = protocol,
        ip = node[1],
        port = node[2],
        masterip = host.ip,
        masterport = port.number
      }
      table.insert(discovered, entry)
    end
  end
  return discovered
end

local function store(servers)
  if not nmap.registry.q3m_servers then
    nmap.registry.q3m_servers = {}
  end
  for _, server in ipairs(servers) do
    table.insert(nmap.registry.q3m_servers, server)
  end
end

local function protocols()
  local filter = {}
  local count = {}
  for _, advert in ipairs(nmap.registry.q3m_servers) do
    local key = stdnse.strjoin(":", {advert.ip, advert.port, advert.protocol})
    if filter[key] == nil then
      if count[advert.protocol] == nil then
        count[advert.protocol] = 0
      end
      count[advert.protocol] = count[advert.protocol] + 1
      filter[key] = true
    end
    local mkey = stdnse.strjoin(":", {advert.masterip, advert.masterport})
  end
  local sortable = {}
  for k, v in pairs(count) do
    table.insert(sortable, {k, v})
  end
  table.sort(sortable, function(a, b) return a[2] > b[2] or (a[2] == b[2] and a[1] > b[1]) end)
  local t = tab.new()
  tab.addrow(t, '#', 'PROTOCOL', 'GAME', 'SERVERS')
  for i, p in ipairs(sortable) do
    local pos = i .. '.'
    local protocol = p[1]
    count = p[2]
    local game = KNOWN_PROTOCOLS[protocol]
    if game == "unknown" then
      game = ""
    end
    tab.addrow(t, pos, protocol, game, count)
  end
  return '\n' .. tab.dump(t)
end

action = function(host, port)
  if SCRIPT_TYPE == "postrule" then
    return protocols()
  end
  local outputlimit = nmap.registry.args[SCRIPT_NAME .. ".outputlimit"]
  if not outputlimit then
    outputlimit = 10
  else
    outputlimit = tonumber(outputlimit)
  end
  if outputlimit < 1 then
    outputlimit = nil
  end
  local servers = scan(host, port, KNOWN_PROTOCOLS)
  store(servers)
  local unique = dropdupes(servers, function(t) return string.format("%s: %s:%d", t.protocol, t.ip, t.port) end)
  local formatted = formatresult(unique, outputlimit, KNOWN_PROTOCOLS)
  if #formatted < 1 then
    return
  end
  local response = {}
  table.insert(response, formatted)
  if outputlimit and outputlimit < #servers then
    table.insert(response, string.format('Only %d/%d shown. Use --script-args %s.outputlimit=-1 to see all.', outputlimit, #servers, SCRIPT_NAME))
  end
  return stdnse.format_output(true, response)
end