---
-- Simple Authentication and Security Layer (SASL).
--
-- The library contains some low level functions and a high level class.
--
-- The Helper class contains the following facilities:
-- * <code>new</code>: This is the SASL object constructor.
-- * <code>set_mechanisms</code>: Constructs the authentication mechanisms
--                                list. Each SASL object must have its own
--                                list of mechanisms to use.
-- * <code>set_callbacks</code>: Associates authentication mechanisms and
--                               authentication functions (callbacks).
-- * <code>next_mechanism</code>: Iterate over the SASL authentication
--                                mechanisms list, and make the next
--                                mechanism on the list the one to use.
-- * <code>select_mechanism</code>: Selects from the SASL authentication
--                                  mechanism list the mechanism to use.
-- * <code>authenticate</code>: Authenticates to the remote server.
-- * <code>reset_callbacks</code>: Resets the callback functions.
-- * <code>reset</code>: Resets the state of the SASL object.
-- @copyright Same as Nmap--See http://nmap.org/book/man-legal.html


module(... or "sasl", package.seeall)

local HAVE_SSL = false

require 'stdnse'
require 'base64'

local MECHANISMS = {
  -- ["mechanism"] = callback function
}

if pcall(require, 'openssl') then
  HAVE_SSL = true
else
  stdnse.print_debug(1,
    "sasl.lua: OpenSSL not present, SASL support limited.")
end

local function sasl_is_success(data)
  return true, data
end

local function sasl_receive(socket)
  local st, ret = socket:receive_lines(1)
  if st then
    st, ret = sasl_is_success(ret)
  end
  return st, ret
end

if HAVE_SSL then
  --- Try to login using the <code>CRAM-MD5<code> mechanism.
  --
  -- @param socket connected to the server.
  -- @param username string.
  -- @param password string.
  -- @return ******
  function cram_md5(socket, username, password)
    local st, ret = socket:send("AUTH CRAM-MD5\r\n")
    return st, ret
  end
  MECHANISMS["CRAM-MD5"] = cram_md5

  --- Try to login using the <code>DIGEST-MD5<code> mechanism.
  --
  -- @param socket connected to the server.
  -- @param username string.
  -- @param password string.
  -- @return ******
  function digest_md5(socket, username, password)
    local st, ret = socket:send("AUTH DIGEST\r\n")
    return st, ret
  end
  MECHANISMS["DIGEST-MD5"] = digest_md5
end

--- Try to login using the <code>PLAIN<code> mechanism.
--
-- @param socket connected to the server.
-- @param username string.
-- @param password string.
-- @return ******
function plain(socket, user, passwd)
  local st, ret = socket:send("AUTH PLAIN\r\n")
  st, ret = socket:send(base64.enc(user.."\0"..user.."\0"..passwd.."\r\n"))
  if st then
    st, ret = sasl_receive(socket)
  end
  return st, ret
end
MECHANISMS["PLAIN"] = plain

--- Returns the supported authentication mechanisms that this library
-- supports and which can be used by the client.
--
-- @param mechanisms The list of the authentication mechanisms
--        supported by the server.
-- @return list The authentication mechanisms list that this library
--         supports and which can be used by the client.
local function get_sasl_mechs(mechanisms)
  local new_mechs = {}
  for _, m in ipairs(mechanisms) do
    local mech = string.upper(m)
    if MECHANISMS[mech] then
      new_mechs[#new_mechs + 1] = {
        mechanism = mech,
        callback = MECHANISMS[mech],
      }
    else
      stdnse.print_debug(3,
          "sasl library does not support '%s' mechanism", m)
    end
  end
  return new_mechs
end

--- This is the SASL Helper class, script writers should use it to create
-- SASL objects.
--
-- Each new SASL object contains:
-- 
-- * The list of the authentication mechanisms associated with it:
--   mechanisms = {
--      { 
--        mechanism = "CRAM-MD5",
--        callback = cram_md5, -- the function that handles cram-md5.
--      },
--      {
--        mechanism = "PLAIN",
--        callback = pain, -- the plain authentication function.
--      },
--      ...
--   }
--   The order of the authentication mechanisms is always respected.
-- 
-- The followings are used to reference the current used mechanisms:
-- * index: Reference the current used mechanism in the previous
--          mechanisms list
-- 
-- * current_mechanism: A string that contains the name of the current
--                      used mechanism.
-- * current_callback: A reference to the authentication function that is
--                     associated with the current authentication
--                     mechanism.
--
-- Usage of the Helper class:
-- local auth_sasl = sasl.Helper.new()
-- auth_sasl:set_mechanisms("PLAIN CRAM-MD5")
-- do 
--    local mech = auth_sasl:get_current_mechanism()
--    -- round 1: mech == "PLAIN", round 2: mech == "CRAM-MD5" ...
--    local st, ret = auth_sasl:authenticate(socket, username, password)
-- until not auth_sasl:next_mechanism()
Helper = {

  --- SASL object constructor
  --
  -- @param mechanisms A list or a string of the authentication
  --        mechanisms to use (optional parameter).
  -- @param callbacks A list of authenticaion mechanisms with their
  --        associated callback functions (optional parameter).
  -- @usage
  -- local auth_sasl = sasl.Helper:new()
  -- local auth_sasl = sasl.Helper:new("CRAM-MD5 PLAIN")
  -- local auth_sasl = sasl.Helper:new({"CRAM-MD5", "PLAIN"})
  -- local auth_sasl = sasl.Helper:new({"CRAM-MD5", "PLAIN"},
  --                                   {["CRAM-MD5"] = my_cram_md5_func})
  -- @return sasl object.
  new = function(self, mechanisms, callbacks)
    local o = {}
    setmetatable(o, self)
    self.__index = self
    
    -- these values must always point to the current
    -- authentication mechanism and its callback
    -- They are automatically set by the set_mechanisms() function
    -- self.idx = 0
    -- self.current_mechanism = ""
    -- self.current_callback = nil

    self:set_mechanisms(mechanisms)
    self:set_callbacks(callbacks)
    
    return o
  end,

  --- Constructs the SASL object authentication mechanisms list.
  --
  -- If this list was not given in the SASL <code>new</code> constructor,
  -- then you must call this methode to associate the mechanisms with
  -- the SASL object, each SASL object must have its own list.
  --
  -- This methode will automatically set the current mechanims to use
  -- to the first one in the provided mechanisms list parameter.
  -- The order of the authentication mechanisms is always respected.
  --
  -- @param mechanisms A list or a string of the authentication
  --        mechanisms to use.
  -- @usage
  -- sasl_auth:set_mechanisms({"PLAIN", "CRAM-MD5"})
  -- sasl_auth:set_mechanisms("PLAIN LOGIN CRAM-MD5 DIGEST-MD5")
  set_mechanisms = function(self, mechanisms)
    local my_mechs = {}
    if mechanisms then
      if type(mechanisms) == "string" and string.len(mechanisms) > 0 then
        self.mechanisms = get_sasl_mechs(stdnse.strsplit(' ', mechanisms))
      elseif type(mechanisms) == "table" and next(mechanisms) then
        self.mechanisms = get_sasl_mechs(mechanisms)
      end
      self.index = 1
      self.current_mechanism = self.mechanisms[self.index].mechanism
      self.current_callback = self.mechanisms[self.index].callback
    else
      self.mechanisms = {}
      self.index = 0
      self.current_mechanism = ""
      self.current_callback = nil
    end
  end,

  --- Associates the authentication mechanisms with their callbacks.
  --
  -- This function will also update the current mechanism callback.
  --
  -- @param callbacks A list of authenticaion mechanisms with their
  --        associated callback functions .
  -- @usage
  -- function cram_md5_handle_func(_, socket, username, password)
  --    -- handle CRAM-MD5 authentication
  -- end
  -- auth_sasl:set_callbacks({["CRAM-MD5"] = cram_md5_handle_func,
  --                          ["PLAIN"] = plain_handle_func,
  --                          ...})
  set_callbacks = function(self, callbacks)
    if callbacks and next(callbacks) then
      if self.mechanisms and next(self.mechanisms) then
        for _, mtable in ipairs(self.mechanisms) do
          if callbacks[mtable.mechanism] then
            mtable.callback = callbacks[mtable.mechanism]
          end
        end
        -- update the current callback function
        self.current_callback = self.mechanisms[self.index].callback
      end
    end
  end,

  --- Resets all the authentication mechanisms functions (callbacks) to
  -- their original state.
  --
  -- This will also reset the current authentication mechanism callback.
  reset_callbacks = function(self)
    if self.mechanisms and next(self.mechanisms) then
      for _, mtable in ipairs(self.mechanisms) do
        mtable.callback = MECHANISMS[mtable.mechanism]
      end
      -- update the current callback function
      self.current_callback = self.mechanisms[self.index].callback
    end
  end,

  --- Makes the next mechanism in the SASL object authentication
  --  mechainsms list the one to use.
  --
  -- @return string The next mechanism name to use on success. This
  --         methode returns nil on failures or if the list is at its end.
  next_mechanism = function(self)
    if self.mechanisms and next(self.mechanisms) then
      if self.mechanisms[self.index + 1] then
        self.index = self.index + 1
        self.current_mechanism = self.mechanisms[self.index].mechanism
        self.current_callback = self.mechanisms[self.index].callback
        return self:get_current_mechanism()
      end
    end
    return nil
  end,

  --- Selects the mechanism to use from the SASL object authentication
  -- mechanisms list.
  -- Returns nil if it failed to set the current mechanism.
  select_mechanism = function(self, mechanism)
    if self.mechanisms and next(self.mechanisms) then
      local idx = 1
      for _, mtable in ipairs(self.mechanisms) do
        local mech = string.upper(mechanism)
        if mtable.mechanism == mech then
          self.index = idx
          self.current_mechanism = self.mechanisms[self.index].mechanism
          self.current_callback = self.mechanisms[self.index].callback
          return self:get_current_mechanism()
        end
        idx = idx + 1
      end
    end
    return nil
  end,

  --- Resets all the data of the SASL object.
  --
  -- This methode will free the internal SASL authentication mechanisms
  -- list.
  reset = function(self)
    self:set_mechanisms()
  end,

  --- Returns a list containing the authentication mechanisms
  -- associated with the current SASL object.
  --
  -- @return list the SASL authentication mechanisms, on failures or
  --         if there are not authentication mechanisms that are
  --         associated with the SASL object this methode will return nil.
  get_mechanisms = function(self)
    if self.mechanisms and next(self.mechanisms) then
      local list = {}
      for _, mtable in ipairs(self.mechanisms) do
        list[#list + 1] = mtable.mechanism
      end
      return list
    end
  end,

  --- Returns a string containing the authentication mechanisms
  -- associated with the current SASL object.
  --
  -- @param separator A string to use to separate the different
  --        mechanisms.
  -- @return string that lists the authentication mechanisms associated
  --         with the current SASL object. On errors or if there are not
  --         available authentication mechanisms this methode will return
  --         nil.
  mechs_to_string = function(self, separator)
    local list = self:get_mechanisms()
    if list then
      return table.concat(list, separator or ' ')
    end
  end,

  --- Returns the current used authentication mechanism
  -- @return string The current used authentication mechanism.
  get_current_mechanism = function(self)
    return self.current_mechanism
  end,

  --- Authenticate to the remote server.
  --
  -- @param socket connected to the server.
  -- @param username to use.
  -- @param password to use.
  -- @return *******
  authenticate = function(self, socket, username, password)
    return self:current_callback(socket, username, password)
  end,
}