local nmap = require "nmap"
local shortport
local stdnse = require "stdnse"
local oops = require "oops"
_ENV = stdnse.module("comm", stdnse.seeall)
local REQUEST_TIMEOUT = 2000
local function get_timeouts(host, opts)
  local connect_timeout, request_timeout
  if opts and opts.connect_timeout then
    connect_timeout = opts.connect_timeout
  elseif opts and opts.timeout then
    connect_timeout = opts.timeout
  else
    connect_timeout = stdnse.get_timeout(host)
  end
  if opts and opts.request_timeout then
    request_timeout = opts.request_timeout
  elseif opts and opts.timeout then
    request_timeout = opts.timeout
  else
    request_timeout = REQUEST_TIMEOUT
  end
  request_timeout = request_timeout + connect_timeout
  return connect_timeout, request_timeout
end
local setup_connect = function(host, port, opts)
  local sock = nmap.new_socket()
  local connect_timeout, request_timeout = get_timeouts(host, opts)
  sock:set_timeout(connect_timeout)
  if type(host) == "string" and opts.any_af then
    local status, addrs = nmap.resolve(host)
    if status then
      host = {ip = addrs[1], targetname = host}
    end
  end
  local status, err = sock:connect(host, port, opts.proto)
  if not status then
    sock:close()
    return oops.raise("Could not connect", status, err)
  end
  sock:set_timeout(request_timeout)
  return true, sock
end
local read = function(sock, opts)
  if opts.lines then
    return oops.raise("receive_lines failed", sock:receive_lines(opts.lines))
  end
  if opts.bytes then
    return oops.raise("receive_bytes failed", sock:receive_bytes(opts.bytes))
  end
  return oops.raise("receive failed", sock:receive())
end
get_banner = function(host, port, opts)
  opts = opts or {}
  opts.recv_before = true
  local socket, errmsg, correct, banner = oops.raise("tryssl failed", tryssl(host, port, nil, opts))
  if socket then
    socket:close()
    return true, banner
  end
  return false, errmsg
end
exchange = function(host, port, data, opts)
  opts = opts or {}
  local status, sock = setup_connect(host, port, opts)
  local ret
  if not status then
    return oops.raise("Failed to connect", status, sock)
  end
  status, ret = sock:send(data)
  if not status then
    sock:close()
    return oops.raise("Failed to send", status, ret)
  end
  status, ret = read(sock, opts)
  sock:close()
  return oops.raise("Faield to read", status, ret)
end
local function is_ssl(port)
  shortport = shortport or require "shortport"
  return shortport.ssl(nil, port)
end
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
  end
  return "tcp","ssl"
end
function opencon(host, port, data, opts)
  opts = opts or {}
  local status, sd = setup_connect(host, port, opts)
  if not status then
    return oops.raise("Failed to connect", false, sd)
  end
  local response, early_resp
  if opts.recv_before then status, early_resp = oops.raise("read failed", read(sd, opts)) end
  if data and #data > 0 then
    sd:send(data)
    status, response = oops.raise("receive failed", sd:receive())
  else
    response = early_resp
  end
  if not status then
    sd:close()
  end
  return status and sd, response, early_resp
end
function tryssl(host, port, data, opts)
  opts = opts or {}
  if not data and not opts.recv_before then
    stdnse.debug1(
      "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")
  end
  opts.proto = opt1
  local sd, response, early_resp = oops.raise(("%s failed"):format(opt1), opencon(host, port, data, opts))
  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
local unittest = require "unittest"
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")
return _ENV;
