Make comm.lua able to talk DTLS

nselib/comm.lua

@@ -25,6 +25,7 @@
 local nmap = require "nmap"
 local shortport
 local stdnse = require "stdnse"
+local tableaux = require "tableaux"
 local oops = require "oops"
 _ENV = stdnse.module("comm", stdnse.seeall)
 
@@ -60,7 +61,10 @@ end
 
 -- Sets up the socket and connects to host:port
 local setup_connect = function(host, port, opts)
-  local sock = nmap.new_socket()
+  local sock = nmap.new_socket(
+    (opts.proto ~= "ssl" and opts.proto)
+    or (type(port) == "table" and port.protocol)
+    or nil)
 
   local connect_timeout, request_timeout = get_timeouts(host, opts)
 
@@ -174,25 +178,12 @@ end
 -- @return Best option ("tcp" or "ssl")
 -- @return Worst option ("tcp" or "ssl")
 local function bestoption(port)
-  if type(port) == 'table' then
-    if port.protocol == "udp" then
-      stdnse.debug2("DTLS (SSL over UDP) is not supported")
-      return "udp", "udp"
-    end
-    if port.version and port.version.service_tunnel and port.version.service_tunnel == "ssl" then return "ssl","tcp" end
-    if port.version and port.version.name_confidence and port.version.name_confidence > 6 then return "tcp","ssl" end
-    local _port = {
-      number = port.number,
-      service = port.service,
-      protocol = port.protocol or "tcp",
-      state = port.state or "open",
-      version = port.version
-    }
-    if is_ssl(_port) then return "ssl","tcp" end
-  elseif type(port) == 'number' then
-    if is_ssl({number=port, protocol="tcp", state="open"}) then return "ssl","tcp" end
+  assert(type(port) == 'table', "bestoption: port must be a table")
+  assert(port.protocol, "bestoption: port table must have protocol field")
+  if is_ssl(port) then
+    return "ssl", port.protocol
   end
-  return "tcp","ssl"
+  return port.protocol, "ssl"
 end
 
 --- This function opens a connection, sends the first data payload and
@@ -263,26 +254,42 @@ end
 -- of the first receive (before sending data)
 function tryssl(host, port, data, opts)
   opts = opts or {}
-  if not data and not opts.recv_before then
-    stdnse.debug1(
+  assert(opts.proto ~= "ssl", "tryssl: opts.proto must not be 'ssl'")
+  local our_port
+  if type(port) == 'table' then
+    if (opts.proto) then
+      assert(opts.proto == port.protocol, "tryssl: opts.proto mismatch port.protocol")
+    end
+    our_port = tableaux.tcopy(port)
+    our_port.state = "open"
+  else
+    our_port = {
+      number = port,
+      protocol = opts.proto or "tcp",
+      state = "open",
+    }
+  end
+  if not data then
+    assert(our_port.protocol ~= "udp",
+      "Using comm.tryssl with UDP requires first data payload.\n\z
+      Impossible to test the connection for the correct protocol!"
+      )
+    assert(opts.recv_before,
       "Using comm.tryssl without either first data payload or opts.recv_before.\n\z
       Impossible to test the connection for the correct protocol!"
       )
   end
-  local opt1, opt2 = bestoption(port)
-  local best = opt1
-  if opts.proto=="udp" then
-    stdnse.debug2("DTLS (SSL over UDP) is not supported")
+  local best = "none"
+  local sd, response, early_resp
+  for _, proto in { bestoption(our_port) } do
+    opts.proto = proto
+    sd, response, early_resp = oops.raise(("%s failed"):format(proto),
+      opencon(host, our_port, data, opts))
+    if sd then
+      best = proto
+      break
+    end
   end
-  opts.proto = opt1
-  local sd, response, early_resp = oops.raise(("%s failed"):format(opt1), opencon(host, port, data, opts))
-  -- Try the second option (If udp, then both options are the same; skip it)
-  if not sd and opt1 ~= "udp" then
-    opts.proto = opt2
-    sd, response, early_resp = oops.raise(("%s failed"):format(opt2), opencon(host, port, data, opts))
-    best = opt2
-  end
-  if not sd then best = "none" end
   return sd, response, best, early_resp
 end
 
@@ -291,8 +298,6 @@ if not unittest.testing() then
   return _ENV
 end
 test_suite = unittest.TestSuite:new()
-test_suite:add_test(unittest.table_equal({bestoption(443)}, {"ssl", "tcp"}), "bestoption ssl number")
-test_suite:add_test(unittest.table_equal({bestoption(80)}, {"tcp", "ssl"}), "bestoption tcp number")
 test_suite:add_test(unittest.table_equal({bestoption({number=8443,protocol="tcp",state="open"})}, {"ssl", "tcp"}), "bestoption ssl table")
 test_suite:add_test(unittest.table_equal({bestoption({number=1234,protocol="tcp",state="open"})}, {"tcp", "ssl"}), "bestoption tcp table")