local base64 = require "base64"
local comm = require "comm"
local sasl = require "sasl"
local stdnse = require "stdnse"
local string = require "string"
local stringaux = require "stringaux"
local table = require "table"
_ENV = stdnse.module("smtp", stdnse.seeall)
local ERROR_MESSAGES = {
  ["EOF"]     = "connection closed",
  ["TIMEOUT"] = "connection timeout",
  ["ERROR"]   = "failed to receive data"
}
local SMTP_CMD = {
  ["EHLO"] = {
    cmd = "EHLO",
    success = {
      [250] = "Requested mail action okay, completed",
    },
    errors = {
      [421] = "<domain> Service not available, closing transmission channel",
      [500] = "Syntax error, command unrecognised",
      [501] = "Syntax error in parameters or arguments",
      [504] = "Command parameter not implemented",
      [550] = "Not implemented",
    },
  },
  ["HELP"] = {
    cmd = "HELP",
    success = {
      [211] = "System status, or system help reply",
      [214] = "Help message",
    },
    errors = {
      [500] = "Syntax error, command unrecognised",
      [501] = "Syntax error in parameters or arguments",
      [502] = "Command not implemented",
      [504] = "Command parameter not implemented",
      [421] = "<domain> Service not available, closing transmission channel",
    },
  },
  ["AUTH"] = {
    cmd = "AUTH",
    success = {[334] = ""},
    errors = {
      [501] = "Authentication aborted",
    },
  },
  ["MAIL"] = {
    cmd = "MAIL",
    success = {
      [250] = "Requested mail action okay, completed",
    },
    errors = {
      [451] = "Requested action aborted: local error in processing",
      [452] = "Requested action not taken: insufficient system storage",
      [500] = "Syntax error, command unrecognised",
      [501] = "Syntax error in parameters or arguments",
      [421] = "<domain> Service not available, closing transmission channel",
      [552] = "Requested mail action aborted: exceeded storage allocation",
    },
  },
  ["RCPT"] = {
    cmd = "RCPT",
    success = {
      [250] = "Requested mail action okay, completed",
      [251] = "User not local; will forward to <forward-path>",
    },
    errors = {
      [450] = "Requested mail action not taken: mailbox unavailable",
      [451] = "Requested action aborted: local error in processing",
      [452] = "Requested action not taken: insufficient system storage",
      [500] = "Syntax error, command unrecognised",
      [501] = "Syntax error in parameters or arguments",
      [503] = "Bad sequence of commands",
      [521] = "<domain> does not accept mail [rfc1846]",
      [421] = "<domain> Service not available, closing transmission channel",
    },
  },
  ["DATA"] = {
    cmd = "DATA",
    success = {
      [250] = "Requested mail action okay, completed",
      [354] = "Start mail input; end with <CRLF>.<CRLF>",
    },
    errors = {
      [451] = "Requested action aborted: local error in processing",
      [554] = "Transaction failed",
      [500] = "Syntax error, command unrecognised",
      [501] = "Syntax error in parameters or arguments",
      [503] = "Bad sequence of commands",
      [421] = "<domain> Service not available, closing transmission channel",
      [552] = "Requested mail action aborted: exceeded storage allocation",
      [554] = "Transaction failed",
      [451] = "Requested action aborted: local error in processing",
      [452] = "Requested action not taken: insufficient system storage",
    },
  },
  ["STARTTLS"] = {
    cmd = "STARTTLS",
    success = {
      [220] = "Ready to start TLS"
    },
    errors = {
      [501] = "Syntax error (no parameters allowed)",
      [454] = "TLS not available due to temporary reason",
    },
  },
  ["RSET"] = {
    cmd = "RSET",
    success = {
      [200] = "nonstandard success response, see rfc876)",
      [250] = "Requested mail action okay, completed",
    },
    errors = {
      [500] = "Syntax error, command unrecognised",
      [501] = "Syntax error in parameters or arguments",
      [504] = "Command parameter not implemented",
      [421] = "<domain> Service not available, closing transmission channel",
    },
  },
  ["VRFY"] = {
    cmd = "VRFY",
    success = {
      [250] = "Requested mail action okay, completed",
      [251] = "User not local; will forward to <forward-path>",
    },
    errors = {
      [500] = "Syntax error, command unrecognised",
      [501] = "Syntax error in parameters or arguments",
      [502] = "Command not implemented",
      [504] = "Command parameter not implemented",
      [550] = "Requested action not taken: mailbox unavailable",
      [551] = "User not local; please try <forward-path>",
      [553] = "Requested action not taken: mailbox name not allowed",
      [421] = "<domain> Service not available, closing transmission channel",
    },
  },
  ["EXPN"] = {
    cmd = "EXPN",
    success = {
      [250] = "Requested mail action okay, completed",
    },
    errors = {
      [550] = "Requested action not taken: mailbox unavailable",
      [500] = "Syntax error, command unrecognised",
      [501] = "Syntax error in parameters or arguments",
      [502] = "Command not implemented",
      [504] = "Command parameter not implemented",
      [421] = "<domain> Service not available, closing transmission channel",
    },
  },
}
get_domain = function(host)
  local nmap_domain = "nmap.scanme.org"
  local result = stdnse.get_script_args("smtp.domain")
  if not result then
    if type(host) == "table" then
      if host.targetname then
        result = host.targetname
      elseif (host.name and #host.name ~= 0) then
        result = host.name
      end
    end
  end
  return result or nmap_domain
end
get_auth_mech = function(response)
  local list = {}
  for _, line in pairs(stringaux.strsplit("\r?\n", response)) do
    local authstr = line:match("%d+%-AUTH%s(.*)$")
    if authstr then
      for mech in authstr:gmatch("[^%s]+") do
        table.insert(list, mech)
      end
      return list
    end
  end
  return nil
end
check_reply = function(cmd, reply)
  local code, msg = string.match(reply, "^([0-9]+)%s*")
  if code then
    cmd = cmd:upper()
    code = tonumber(code)
    if SMTP_CMD[cmd] then
      if SMTP_CMD[cmd].success[code] then
        return true, reply
      end
    else
      stdnse.debug3(
          "SMTP: check_smtp_reply failed: %s not supported", cmd)
      return false, string.format("SMTP: %s %s", cmd, reply)
    end
  end
  stdnse.debug3(
      "SMTP: check_smtp_reply failed: %s %s", cmd, reply)
  return false, string.format("SMTP: %s %s", cmd, reply)
end
query = function(socket, cmd, data, lines)
  if data then
    cmd = cmd.." "..data
  end
  local st, ret = socket:send(string.format("%s\r\n", cmd))
  if not st then
    socket:close()
    stdnse.debug3("SMTP: failed to send %s request.", cmd)
    return st, string.format("SMTP failed to send %s request.", cmd)
  end
  st, ret = socket:receive_lines(lines or 1)
  if not st then
    socket:close()
    stdnse.debug3("SMTP %s: failed to receive data: %s.",
                    cmd, (ERROR_MESSAGES[ret] or 'unspecified error'))
    return st, string.format("SMTP %s: failed to receive data: %s",
                    cmd, (ERROR_MESSAGES[ret] or 'unspecified error'))
  end
  return st, ret
end
connect = function(host, port, opts)
  local socket, _, ret
  if opts.ssl then
    socket, _, _, ret = comm.tryssl(host, port, '', opts)
  else
    socket, _, ret = comm.opencon(host, port, nil, opts)
  end
  if not socket then
    return socket, (ERROR_MESSAGES[ret] or 'unspecified error')
  end
  return socket, ret
end
starttls = function(socket)
  local st, reply, ret
  st, reply = query(socket, "STARTTLS")
  if not st then
    return st, reply
  end
  st, ret = check_reply('STARTTLS', reply)
  if not st then
    quit(socket)
    return st, ret
  end
  st, ret = socket:reconnect_ssl()
  if not st then
    socket:close()
    return st, ret
  end
  return true, reply
end
ehlo = function(socket, domain)
  local st, ret, response
  st, response = query(socket, "EHLO", domain)
  if not st then
    return st, response
  end
  st, ret = check_reply("EHLO", response)
  if not st then
    quit(socket)
    return st, ret
  end
  return st, response
end
help = function(socket)
  local st, ret, response
  st, response = query(socket, "HELP")
  if not st then
    return st, response
  end
  st, ret = check_reply("HELP", response)
  if not st then
    quit(socket)
    return st, ret
  end
  return st, response
end
mail = function(socket, address, esmtp_opts)
  local st, ret, response
  if esmtp_opts and next(esmtp_opts) then
    local data = ""
    for k,v in pairs(esmtp_opts) do
      k = k:upper()
      data = string.format("%s %s=%s", data, k, v)
    end
    st, response = query(socket, "MAIL",
                      string.format("FROM:<%s>%s",
                      address, data))
  else
    st, response = query(socket, "MAIL",
                      string.format("FROM:<%s>", address))
  end
  if not st then
    return st, response
  end
  st, ret = check_reply("MAIL", response)
  if not st then
    quit(socket)
    return st, ret
  end
  return st, response
end
recipient = function(socket, address)
  local st, ret, response
  st, response = query(socket, "RCPT",
                      string.format("TO:<%s>", address))
  if not st then
    return st, response
  end
  st, ret = check_reply("RCPT", response)
  if not st then
    quit(socket)
    return st, ret
  end
  return st, response
end
datasend = function(socket, data)
  local st, ret, response
  st, response = query(socket, "DATA")
  if not st then
    return st, response
  end
  st, ret = check_reply("DATA", response)
  if not st then
    quit(socket)
    return st, ret
  end
  if data then
    st, response = query(socket, data.."\r\n.")
    if not st then
      return st, response
    end
    st, ret = check_reply("DATA", response)
    if not st then
      quit(socket)
      return st, ret
    end
  end
  return st, response
end
reset = function(socket)
  local st, ret, response
  st, response = query(socket, "RSET")
  if not st then
    return st, response
  end
  st, ret = check_reply("RSET", response)
  if not st then
    quit(socket)
    return st, ret
  end
  return st, response
end
verify = function(socket, mailbox)
  local st, ret, response
  st, response = query(socket, "VRFY", mailbox)
  st, ret = check_reply("VRFY", response)
  if not st then
    quit(socket)
    return st, ret
  end
  return st, response
end
quit = function(socket)
  stdnse.debug3("SMTP: sending 'QUIT'.")
  socket:send("QUIT\r\n")
  socket:close()
end
login = function(socket, username, password, mech)
  assert(mech == "LOGIN" or mech == "PLAIN" or mech == "CRAM-MD5"
    or mech == "DIGEST-MD5" or mech == "NTLM",
    ("Unsupported authentication mechanism (%s)"):format(mech or "nil"))
  local status, response = query(socket, "AUTH", mech)
  if ( not(status) ) then
    return false, "ERROR: Failed to send AUTH to server"
  end
  if ( mech == "LOGIN" ) then
    local tmp = response:match("334 (.*)")
    if ( not(tmp) ) then
      return false, "ERROR: Failed to decode LOGIN response"
    end
    tmp = base64.dec(tmp):lower()
    if ( not(tmp:match("^username")) ) then
      return false, ("ERROR: Expected \"Username\", but received (%s)"):format(tmp)
    end
    status, response = query(socket, base64.enc(username))
    if ( not(status) ) then
      return false, "ERROR: Failed to read LOGIN response"
    end
    tmp = response:match("334 (.*)")
    if ( not(tmp) ) then
      return false, "ERROR: Failed to decode LOGIN response"
    end
    tmp = base64.dec(tmp):lower()
    if ( not(tmp:match("^password")) ) then
      return false, ("ERROR: Expected \"password\", but received (%s)"):format(tmp)
    end
    status, response = query(socket, base64.enc(password))
    if ( not(status) ) then
      return false, "ERROR: Failed to read LOGIN response"
    end
    if ( response:match("^235") ) then
      return true, "Login success"
    end
    return false, response
  end
  if ( mech == "NTLM" ) then
    status, response = query(socket, "TlRMTVNTUAABAAAAB7IIogYABgA3AAAADwAPACgAAAAFASgKAAAAD0FCVVNFLUFJUi5MT0NBTERPTUFJTg==")
    if ( not(status) ) then return false, "ERROR: Failed to receive NTLM challenge" end
  end
  local chall = response:match("^334 (.*)")
  chall = (chall and base64.dec(chall))
  if (not(chall)) then return false, "ERROR: Failed to retrieve challenge" end
  local mech_params = { username, password, chall, "smtp" }
  local auth_data = sasl.Helper:new(mech):encode(table.unpack(mech_params))
  auth_data = base64.enc(auth_data)
  status, response = query(socket, auth_data)
  if ( not(status) ) then
    return false, ("ERROR: Failed to authenticate using SASL %s"):format(mech)
  end
  if ( mech == "DIGEST-MD5" ) then
    local rspauth = response:match("^334 (.*)")
    if ( rspauth ) then
      rspauth = base64.dec(rspauth)
      status, response = query(socket,"")
    end
  end
  if ( response:match("^235") ) then return true, "Login success" end
  return false, response
end
return _ENV;
