-- -*- mode: lua; lua-indent-level: 2; ispell-local-dictionary: "british" -*-

local os = require "os"
local bin = require "bin"
local bit = require "bit"
local comm = require "comm"
local math = require "math"
local nmap = require "nmap"
local stdnse = require "stdnse"
local string = require "string"
local table = require "table"
local shortport = require "shortport"
local nsedebug = require "nsedebug"


-- XXX: Add option to try other probe pl compositions on responding
-- targets.


description = [[
Probes for Simple service discovery protocol (ssdp) servers.
]]

---
-- @usage
-- nmap -n -sU -Pn --script ssdp -pU:1900 --script-args ssdp.mx=1s -- <target>
--
-- @args
-- ssdp.mx timeout set plus half a second on the socket and in the
-- ssdp m-search request after flooring to integer, should be between
-- 1 and 5 seconds inclusive according to
-- upnp-arch-devicearchitecture-v1.1 (default = 1s).
--
-- ssdp.drop-host-header if supplied send an ssdp m-search probe
-- without the host header required by
-- upnp-arch-devicearchitecture-v1.1; some implementations answer
-- anyway providing greater amplification.
--
-- @output
-- PORT     STATE SERVICE
-- 1900/udp open  ssdp
-- | ssdp: 
-- |   product: RAIDiator/4.1 DLNADOC/1.50 UPnP/1.0 MiniDLNA/1.0
-- |   matched targets
-- |     uuid:d02cfb64-717d-11e4-b5f1-0f89aa5e4518
-- |     upnp:rootdevice
-- |     urn:schemas-upnp-org:device:MediaServer:1
-- |     urn:schemas-upnp-org:service:ContentDirectory:1
-- |     urn:schemas-upnp-org:service:ConnectionManager:1
-- |     urn:microsoft.com:service:X_MS_MediaReceiverRegistrar:1
-- |   packet ratio: 6
-- |_  payload ratio: 2114/90=23.5
--
-- @xmloutput
-- <port protocol="udp" portid="1900">
--   <state state="open" reason="script-set" reason_ttl="0"/>
--   <service name="ssdp"
--            product="RAIDiator/4.1 DLNADOC/1.50 UPnP/1.0 MiniDLNA/1.0"
--            method="probed" conf="10"/>
--   <script id="ssdp" output="see @output ...">
--     <elem key="product">RAIDiator/4.1 DLNADOC/1.50 UPnP/1.0 MiniDLNA/1.0</elem>
--     <table key="matched_targets">
--       <elem>uuid:d02cfb64-717d-11e4-b5f1-0f89aa5e4518</elem>
--       <elem>upnp:rootdevice</elem>
--       <elem>urn:schemas-upnp-org:device:MediaServer:1</elem>
--       <elem>urn:schemas-upnp-org:service:ContentDirectory:1</elem>
--       <elem>urn:schemas-upnp-org:service:ConnectionManager:1</elem>
--       <elem>urn:microsoft.com:service:X_MS_MediaReceiverRegistrar:1</elem>
--     </table>
--     <elem key="packet_ratio">6</elem>
--     <elem key="payload_ratio">2114/90=23.5</elem>
--   </script>
-- </port>


categories = {"default", "discovery", "safe", "version"}
author = "Ulrik Haugen"
copyright = "Linköpings universitet 2014, Ulrik Haugen 2014"
license = "Same as Nmap--See http://nmap.org/book/man-legal.html"


portrule = shortport.port_or_service(1900, 'ssdp', 'udp')


--- Like assert but put /message/ in the ERROR key in /results_table/ to
-- better suit collate_results and pass 0 as level to error to ensure
-- the error message will not be prefixed with file and line number.
-- /results_table/ may be left out.
local function assert_w_table(condition, message, results_table)
  if condition then
    return condition
  else
    results_table = results_table or {}
    results_table.ERROR = message
    error(results_table, 0)
  end
end


--- Send an ssdp probe and look for responses.
local function probe_server(host, port, mx, drop_host_header)
  local sock = nmap.new_socket()
  sock:set_timeout(mx * 1000 + 500)
  local status, ret = sock:connect(host, port, port.protocol)
  assert_w_table(status, ret)

  local server_info = stdnse.output_table()

  -- Send a multicast m-search request (the mx header does not occur
  -- in unicast m-search) as it seems to be more widely recognised,
  -- over unicast. If you wanted to boost amplification without
  -- dropping the host header you could substitute host.ip for the
  -- multicast address as it can be shorter but never longer.
  local probe_req = {
    'M-SEARCH * HTTP/1.1',
    string.format('HOST:239.255.255.250:%s', port.number),
    'MAN:"ssdp:discover"',
    string.format('MX:%s', math.floor(mx)),
    'ST:ssdp:all',
    '',
    '', }
  if drop_host_header then
    table.remove(probe_req, 2)
  end
  local probe_pl = table.concat(probe_req, '\r\n')
  -- nsedebug.print_hex(probe_pl)
  status, ret = sock:send(probe_pl)
  if not status then
    sock:close()
    assert_w_table(false, ret)
  end
  local res_pls = {}
  while true do
    status, ret = sock:receive_bytes(1)
    -- nsedebug.print_hex(ret)
    if not status then
      break
    end
    table.insert(res_pls, ret)
  end
  sock:close()
  if #res_pls == 0 then
    assert_w_table(false, string.format("No response to probe: %s", ret))
  end

  nmap.set_port_state(host, port, 'open')

  local idx, res_pl
  local res_pls_len = 0
  server_info.matched_targets = {}
  for idx, res_pl in pairs(res_pls) do
    if not server_info.product and string.match(res_pl, '^HTTP/1.1 200 OK') then
      port.version.name = 'ssdp'
      port.version.product = string.match(res_pl, '\r\nSERVER: *([%g ]+)')
      nmap.set_port_version(host, port)
      server_info.product = port.version.product
    end

    res_pls_len = res_pls_len + res_pl:len()
    local matched_target = string.match(res_pl, '\r\nST: *([%g ]+)')
    if matched_target then
      table.insert(server_info.matched_targets, matched_target)
    end
  end

  server_info.packet_ratio = string.format("%d", #res_pls)
  server_info.payload_ratio = string.format("%d/%d=%.1f",
                                            res_pls_len, probe_pl:len(),
                                            res_pls_len/probe_pl:len())
  return server_info
end


--- Return a function from structured to unstructured output indenting
-- nested tables /offset/ or two spaces with special treatment of name
-- keys and optionally using /xlate_key/ to format keys.
local function make_formatter(offset, xlate_key)
  offset = offset or 2
  xlate_key = xlate_key or function(key) return key:gsub("_", " ") end
  local function select_string(bool, yes, no)
    return ({ [ true ] = yes, [ false ] = no })[bool]
  end

  --- Format /results_table/ as a string starting /indent/ or zero
  -- steps from the margin for the name key and adding offset steps
  -- for other table contents and again for the contents of nested
  -- tables.
  local function formatter(results_table, indent)
    indent = indent or 0
    local output = {}

    if results_table.name then
      table.insert(output,
                   string.format("%s%s",
                                 select_string(indent == 0, ": ", "\n"),
                                 results_table.name))
    end

    for key, value in pairs(results_table) do
      -- name is printed already
      if key ~= 'name' then
        table.insert(output,
                     string.format("\n%s",
                                   string.rep(" ", indent + offset)))
        if type(key) ~= 'number' then
          table.insert(output,
                       string.format("%s%s",
                                     xlate_key(key),
                                     select_string(type(value) == 'table',
                                                   "", ": ")))
        end

        if type(value) == 'table' then
          table.insert(output, formatter(value, indent + offset))
        else
          table.insert(output, value)
        end
      end
    end
    return table.concat(output, '')
  end

  return formatter
end


--- Use /formatter/ to produce unstructured output from
-- /results_table/ considering /status/. Return structured and
-- unstructured output.
local function collate_results(formatter, status, results_table)
  if not status and nmap.debugging() < 1 then
    return nil
  end
  return results_table, formatter(results_table)
end


--- Nmap entry point.
function action(host, port)
  local mx = stdnse.parse_timespec(stdnse.get_script_args(
                                     SCRIPT_NAME .. ".mx")) or 1
  local drop_host_header = stdnse.get_script_args(
    SCRIPT_NAME .. ".drop-host-header")
  return collate_results(make_formatter(),
                         pcall(probe_server, host, port,
                               mx, drop_host_header))
end