---
-- Simple Authentication and Security Layer (SASL).
--
-- The library contains some low level functions and a high level class.
--
-- The <code>DigestMD5</code> class contains all code necessary to calculate
-- a DIGEST-MD5 response based on the servers challenge and the other
-- necessary arguments (@see DigestMD5.new).
-- It can be called throught the SASL helper or directly like this:
-- <code>
-- 	local dmd5 = DigestMD5:new(chall, user, pass, "AUTHENTICATE", nil, "imap")
-- 	local digest = dmd5:calcDigest()
-- </code>
-- 
-- The <code>NTLM</code> class contains all code necessary to calculate a
-- NTLM response based on the servers challenge and the other necessary
-- arguments (@see NTLM.new). It can be called through the SASL helper or
-- directly like this:
-- <code>
-- 	local ntlm = NTLM:new(chall, user, pass)
--  local response = ntlm:calcResponse()
-- </code>
--
-- The Helper class contains the high level methodes:
-- * <code>new</code>: This is the SASL object constructor.
-- * <code>set_mechanism</code>: Sets the authentication mechanism to use.
-- * <code>set_callback</code>: Sets the encoding function to use.
-- * <code>encode</code>: Encodes the parameters according to the
--                        authentication mechanism.
-- * <code>reset_callback</code>: Resets the authentication function.
-- * <code>reset</code>: Resets the SASL object.
--
-- The script writers should use the Helper class to create SASL objects,
-- and they can also use the low level functions to customize their
-- encoding functions.
--
-- @copyright Same as Nmap--See http://nmap.org/book/man-legal.html

-- Version 0.2
-- Created 07/17/2011 - v0.1 - 	Created by Djalal Harouini
-- Revised 07/18/2011 - v0.2 - 	Added NTLM, DIGEST-MD5 classes


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

local HAVE_SSL = false

require 'stdnse'
require 'base64'
require 'smbauth'

local MECHANISMS = { }

-- Calculates a DIGEST MD5 response
DigestMD5 = {
	
	--- Instantiates DigestMD5
	--
	-- @param chall string containing the base64 decoded challenge
	-- @return a new instance of DigestMD5
	new = function(self, chall, username, password, method, uri, service, realm)
		local o = { nc = 0, 
			chall = chall, 
			challnvs = {},
			username = username,
			password = password,
			method = method,
			uri = uri,
			service = service,
			realm = realm }
       	setmetatable(o, self)
        self.__index = self
		o:parseChallenge()
		return o
	end,
	
	-- parses a challenge received from the server
	-- takes care of both quoted and unqoted identifiers
	-- regardless of what RFC says
	parseChallenge = function(self)
		local results = {}
		local start, stop = 0,0
		while(true) do
			local name, value
			start, stop, name = self.chall:find("([^=]*)=%s*", stop + 1)
			if ( not(start) ) then break end
			if ( self.chall:sub(stop + 1, stop + 1) == "\"" ) then
				start, stop, value = self.chall:find("(.-)\"", stop + 2)
			else
				start, stop, value = self.chall:find("([^,]*)", stop + 1)
			end
			self.challnvs[name:lower()] = value
			start, stop = self.chall:find("%s*,%s*", stop + 1)
			if ( not(start) ) then break end
		end
	end,
	
	--- Calculates the digest 
	calcDigest = function( self )
		local uri = self.uri or ("%s/%s"):format(self.service, "localhost")
		local realm = self.realm or self.challnvs.realm or ""
		local cnonce = stdnse.tohex(openssl.rand_bytes( 8 ))
		local qop = "auth"
		self.nc = self.nc + 1
		local A1_part1 = openssl.md5(self.username .. ":" .. (self.challnvs.realm or "") .. ":" .. self.password)
		local A1 = stdnse.tohex(openssl.md5(A1_part1 .. ":" .. self.challnvs.nonce .. ':' .. cnonce))
		local A2 = stdnse.tohex(openssl.md5(("%s:%s"):format(self.method, uri)))
		local digest = stdnse.tohex(openssl.md5(A1 .. ":" .. self.challnvs.nonce .. ":" ..
					("%08d"):format(self.nc)  .. ":" .. cnonce .. ":" ..
					qop .. ":" .. A2))
		
		local response = "username=\"" .. self.username .. "\""
		response = response .. (",%s=\"%s\""):format("realm", realm)
		response = response .. (",%s=\"%s\""):format("nonce", self.challnvs.nonce)
		response = response .. (",%s=\"%s\""):format("cnonce", cnonce)
		response = response .. (",%s=%08d"):format("nc", self.nc)
		response = response .. (",%s=%s"):format("qop", "auth")
		response = response .. (",%s=\"%s\""):format("digest-uri", uri)
		response = response .. (",%s=%s"):format("response", digest)
		response = response .. (",%s=%s"):format("charset", "utf-8")
						
		return response
	end,
	
	
}

-- The NTLM class handling NTLM challenge response authentication
NTLM = {
	
	--- Creates a new instance of the NTLM class
	--
	-- @param chall string containing the challenge received from the server
	-- @param username string containing the username
	-- @param password string containing the password
	-- @return new instance of NTML
	new = function(self, chall, username, password)
		local o = { nc = 0, 
			chall = chall, 
			username = username,
			password = password}
       	setmetatable(o, self)
        self.__index = self
		o:parseChallenge()
		return o
	end,
	
	--- Converst str to "unicode" (adds null bytes for every other byte)
	-- @param str containing string to convert
	-- @return unicode string containing the unicoded str
	to_unicode = function(str)
		local unicode = ""
		for i = 1, #str, 1 do
			unicode = unicode .. bin.pack("<S", string.byte(str, i))
		end
		return unicode
	end,
	
	--- Parses the NTLM challenge as received from the server
	parseChallenge = function(self)
		local NTLM_NegotiateUnicode 			= 0x00000001
		local NTLM_NegotiateExtendedSecurity 	= 0x00080000
		local pos, _, message_type
				
		pos, _, message_type, _, _, 
		_, self.flags, self.chall, _, 
		_, _, _	 = bin.unpack("<A8ISSIIA8LSSI", self.chall)
		
		if ( message_type ~= 0x02 ) then
			error("NTLM parseChallenge expected message type: 0x02")
		end
		
		self.is_extended = ( bit.band(self.flags, NTLM_NegotiateExtendedSecurity) == NTLM_NegotiateExtendedSecurity )
		local is_unicode  = ( bit.band(self.flags, NTLM_NegotiateUnicode) == NTLM_NegotiateUnicode )
		
		self.workstation = "NMAP-HOST"
		self.domain = self.username:match("^(.-)\\(.*)$") or "DOMAIN"
		
		if ( is_unicode ) then
			self.workstation = self.to_unicode(self.workstation)
			self.username = self.to_unicode(self.username)
			self.domain = self.to_unicode(self.domain)
		end	
	end,
	
	--- Calculates the response
	calcResponse = function(self)
		local ntlm, lm = smbauth.get_password_response(nil, self.username, self.domain, self.password, nil, "v1", self.chall, self.is_extended)
		local msg_type = 3
		local response
		local BASE_OFFSET = 72
		local offset
		local encrypted_random_sesskey = ""
		local flags = 0xa2888205  -- (NTLM_NegotiateUnicode | \
		                    		-- NTLM_RequestTarget | \
				                    -- NTLM_NegotiateNTLM | \
				                    -- NTLM_NegotiateAlwaysSign | \
				                    -- NTLM_NegotiateExtendedSecurity | \
				                    -- NTLM_NegotiateTargetInfo | \
				                    -- NTLM_NegotiateVersion | \
				                    -- NTLM_Negotiate128 | \
				                    -- NTLM_Negotiate56)
		
		response = bin.pack("<AI", "NTLMSSP\0", msg_type)

		offset = BASE_OFFSET + #self.workstation + #self.username + #self.domain
		response = response .. bin.pack("<SSI", #lm, #lm, offset)

		offset = offset + #lm
		response = response .. bin.pack("<SSI", #ntlm, #ntlm, offset)

		offset = BASE_OFFSET
		response = response .. bin.pack("<SSI", #self.domain, #self.domain, offset)
		
		offset = BASE_OFFSET + #self.domain
		response = response .. bin.pack("<SSI", #self.username, #self.username, offset)
		
		offset = BASE_OFFSET + #self.domain + #self.username
		response = response .. bin.pack("<SSI", #self.workstation, #self.workstation, offset)
		
		offset = offset + #self.workstation + #lm + #ntlm
		response = response .. bin.pack("<SSI", #encrypted_random_sesskey, #encrypted_random_sesskey, offset)
	
		response = response .. bin.pack("<I", flags)

		-- add version info (major 5, minor 1, build 2600, reserved(1-3) 0,
		-- NTLM Revision 15)
		response = response .. bin.pack("<CCSCCCC", 5, 1, 2600, 0, 0, 0, 15)
		response = response .. self.domain .. self.username .. self.workstation .. ntlm .. lm .. encrypted_random_sesskey
	
		return response
	end
	
}

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

if HAVE_SSL then 
  --- Encodes the parameters using the <code>CRAM-MD5</code> mechanism.
  --
  -- @param username string.
  -- @param password string.
  -- @param challenge The challenge as it is returned by the server.
  -- @return string The encoded string on success, or nil if Nmap was
  --         compiled without OpenSSL.
  function cram_md5_enc(username, password, challenge)
    local encode = stdnse.tohex(openssl.hmac('md5',
                                             password,
                                             challenge))
    return username.." "..encode
  end

  --- Encodes the parameters using the <code>DIGEST-MD5</code> mechanism.
  --
  -- @param username string.
  -- @param password string.
  -- @param challenge The challenge as it is returned by the server.
  -- @param service string containing the service that is requesting the
  --        encryption (eg. POP, IMAP, STMP)
  -- @param uri string containing the URI
  -- @return string The encoded string on success, or nil if Nmap was
  --         compiled without OpenSSL.
  function digest_md5_enc(username, password, challenge, service, uri)
	return DigestMD5:new(challenge, 
							username, 
							password, 
							"AUTHENTICATE", 
							uri, 
							service):calcDigest()
  end

  function ntlm_enc(username, password, challenge)
	return NTLM:new(challenge, username, password):calcResponse()
  end

else
  function cram_md5_enc()
    error("cram_md5_enc not supported without OpenSSL")
  end

  function digest_md5_enc()
    error("digest_md5_enc not supported without OpenSSL")
  end

  function ntlm_enc()
    error("ntlm_enc not supported without OpenSSL")
  end
end

MECHANISMS["CRAM-MD5"] = cram_md5_enc
MECHANISMS["DIGEST-MD5"] = digest_md5_enc
MECHANISMS["NTLM"] = ntlm_enc


--- Encodes the parameters using the <code>PLAIN</code> mechanism.
--
-- @param username string.
-- @param password string.
-- @return string The encoded string.
function plain_enc(username, password)
  return username.."\0"..username.."\0"..password
end
MECHANISMS["PLAIN"] = plain_enc


--- Checks if the given mechanism is supported by this library.
--
-- @param mechanism string to check.
-- @return mechanism if it is supported, otherwise nil.
-- @return callback The mechanism encoding function on success.
function check_mechanism(mechanism)
  local lmech, lcallback
  if mechanism then
    mechanism = string.upper(mechanism)
    if MECHANISMS[mechanism] then
      lmech = mechanism
      lcallback = MECHANISMS[mechanism]
    else
      stdnse.print_debug(3,
        "sasl library does not support '%s' mechanism", mechanism)
    end
  end
  return lmech, lcallback
end

--- This is the SASL Helper class, script writers should use it to create
-- SASL objects.
--
-- Usage of the Helper class:
--
-- local sasl_enc = sasl.Helper.new("CRAM-MD5")
-- local result = sasl_enc:encode(username, password, challenge)
--
-- sasl_enc:set_mechanism("LOGIN")
-- local user, pass = sasl_enc:encode(username, password)
Helper = {

  --- SASL object constructor.
  --
  -- @param mechanism The authentication mechanism to use
  --        (optional parameter).
  -- @param callback The encoding function associated with the
  --        mechanism (optional parameter).
  -- @usage
  -- local sasl_enc = sasl.Helper:new()
  -- local sasl_enc = sasl.Helper:new("CRAM-MD5")
  -- local sasl_enc = sasl.Helper:new("CRAM-MD5", my_cram_md5_func)
  -- @return sasl object.
  new = function(self, mechanism, callback)
    local o = {}
    setmetatable(o, self)
    self.__index = self
    if self:set_mechanism(mechanism) then
      self:set_callback(callback)
    end
    return o
  end,

  --- Sets the SASL mechanism to use.
  --
  -- @param string The authentication mechanism.
  -- @usage
  -- local sasl_enc = sasl.Helper:new()
  -- sasl_enc:set_mechanism("CRAM-MD5")
  -- sasl_enc:set_mechanism("PLAIN")
  -- @return mechanism on success, or nil if the mechanism is not
  --         supported.
  set_mechanism = function(self, mechanism)
    self.mechanism, self.callback = check_mechanism(mechanism)
    return self.mechanism
  end,

  --- Associates A custom encoding function with the authentication
  --  mechanism.
  --
  -- Note that the SASL object by default will have its own
  -- callback functions.
  --
  -- @param callback The function associated with the authentication
  --        mechanism.
  -- @usage
  -- -- My personal CRAM-MD5 encode function
  -- function cram_md5_encode_func(username, password, challenge)
  --    ...
  -- end
  -- local sasl_enc = sasl.Helper:new("CRAM-MD5")
  -- sasl_enc:set_callback(cram_md5_handle_func)
  -- local result = sasl_enc:encode(username, password, challenge)
  set_callback = function(self, callback)
    if callback then
      self.callback = callback
    end
  end,

  --- Resets the encoding function to the default SASL
  --  callback function.
  reset_callback = function(self)
    self.callback = MECHANISMS[self.mechanism]
  end,

  --- Resets all the data of the SASL object.
  --
  -- This methode will clear the specified SASL mechanism.
  reset = function(self)
    self:set_mechanism()
  end,

  --- Returns the current authentication mechanism.
  -- 
  -- @return mechanism on success, or nil on failures.
  get_mechanism = function(self)
    return self.mechanism
  end,

  --- Encodes the parameters according to the specified mechanism.
  --
  -- @param ... The parameters to encode.
  -- @usage
  -- local sasl_enc = sasl.Helper:new("CRAM-MD5")
  -- local result = sasl_enc:encode(username, password, challenge)
  -- local sasl_enc = sasl.Helper:new("PLAIN")
  -- local result = sasl_enc:encode(username, password)
  -- @return string The encoded string on success, or nil on failures.
  encode = function(self, ...)
    return self.callback(...)
  end,
}