local smbauth = require "smbauth"
local stdnse = require "stdnse"
local string = require "string"
local unicode = require "unicode"
_ENV = stdnse.module("sasl", stdnse.seeall)
local HAVE_SSL, openssl = pcall(require, 'openssl')
if ( not(HAVE_SSL) ) then
  stdnse.debug1(
  "sasl.lua: OpenSSL not present, SASL support limited.")
end
local MECHANISMS = { }
if HAVE_SSL then
  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,
    parseChallenge = function(self)
      local results = {}
      if self.chall then
        local start, stop = self.chall:find("^[Dd][Ii][Gg][Ee][Ss][Tt]%s+")
        stop = stop or 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
          name = name:lower()
          self.challnvs[name] = value
          start, stop = self.chall:find("%s*,%s*", stop + 1)
          if ( not(start) ) then break end
        end
      end
    end,
    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"
      local qop_not_specified
      if self.challnvs.qop then
        qop_not_specified = false
      else
        qop_not_specified = true
      end
      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 b1
      if not self.challnvs.algorithm or self.challnvs.algorithm:upper() == "MD5" then
        b1 = stdnse.tohex(openssl.md5(self.username..":"..(self.challnvs.realm or "")..":"..self.password))
      else
        b1 = A1
      end
      local digest_http
      if not qop_not_specified then
        digest_http =  stdnse.tohex(openssl.md5(b1 .. ":" .. self.challnvs.nonce .. ":" ..
          ("%08d"):format(self.nc)  .. ":" .. cnonce .. ":" .. qop .. ":" .. A2))
      else
        digest_http =  stdnse.tohex(openssl.md5(b1 .. ":" .. self.challnvs.nonce .. ":" .. A2))
      end
      local response = "username=\"" .. self.username .. "\""
      .. (",%s=\"%s\""):format("realm", realm)
      .. (",%s=\"%s\""):format("nonce", self.challnvs.nonce)
      .. (",%s=\"%s\""):format("cnonce", cnonce)
      .. (",%s=%08d"):format("nc", self.nc)
      .. (",%s=%s"):format("qop", "auth")
      .. (",%s=\"%s\""):format("digest-uri", uri)
      .. (",%s=%s"):format("response", digest)
      .. (",%s=%s"):format("charset", "utf-8")
      local response_table = {
        username = self.username,
        realm = realm,
        nonce = self.challnvs.nonce,
        cnonce = cnonce,
        nc = ("%08d"):format(self.nc),
        qop = qop,
        ["digest-uri"] = uri,
        algorithm = self.challnvs.algorithm,
        response = digest_http
      }
      return response, response_table
    end,
  }
  NTLM = {
    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,
    parseChallenge = function(self)
      local NTLM_NegotiateUnicode = 0x00000001
      local NTLM_NegotiateExtendedSecurity = 0x00080000
      local pos, _, message_type
      _, message_type, _, _,
      _, self.flags, self.chall, _,
      _, _, _, pos = string.unpack("<c8 I4 I2 I2 I4 I4 c8 I8 I2 I2 I4", self.chall)
      if ( message_type ~= 0x02 ) then
        error("NTLM parseChallenge expected message type: 0x02")
      end
      self.is_extended = ( (self.flags & NTLM_NegotiateExtendedSecurity) == NTLM_NegotiateExtendedSecurity )
      local is_unicode  = ( (self.flags & NTLM_NegotiateUnicode) == NTLM_NegotiateUnicode )
      self.workstation = "NMAP-HOST"
      self.domain = self.username:match("^(.-)\\(.*)$") or "DOMAIN"
      if ( is_unicode ) then
        self.workstation = unicode.utf8to16(self.workstation)
        self.username = unicode.utf8to16(self.username)
        self.domain = unicode.utf8to16(self.domain)
      end
    end,
    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 | \
      response = string.pack("<zI4", "NTLMSSP", msg_type)
      offset = BASE_OFFSET + #self.workstation + #self.username + #self.domain
      response = response .. string.pack("<I2I2I4", #lm, #lm, offset)
      offset = offset + #lm
      response = response .. string.pack("<I2I2I4", #ntlm, #ntlm, offset)
      offset = BASE_OFFSET
      response = response .. string.pack("<I2I2I4", #self.domain, #self.domain, offset)
      offset = BASE_OFFSET + #self.domain
      response = response .. string.pack("<I2I2I4", #self.username, #self.username, offset)
      offset = BASE_OFFSET + #self.domain + #self.username
      response = response .. string.pack("<I2I2I4", #self.workstation, #self.workstation, offset)
      offset = offset + #self.workstation + #lm + #ntlm
      response = response .. string.pack("<I2I2I4", #encrypted_random_sesskey, #encrypted_random_sesskey, offset)
      response = response .. string.pack("<I4", flags)
      response = response .. string.pack("<BBI2 BBBB", 5, 1, 2600, 0, 0, 0, 15)
      response = response .. self.domain .. self.username .. self.workstation .. ntlm .. lm .. encrypted_random_sesskey
      return response
    end
  }
  function cram_md5_enc(username, password, challenge)
    local encode = stdnse.tohex(openssl.hmac('md5',
      password,
      challenge))
    return username.." "..encode
  end
  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
function plain_enc(username, password)
  return username.."\0"..username.."\0"..password
end
MECHANISMS["PLAIN"] = plain_enc
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.debug3(
        "sasl library does not support '%s' mechanism", mechanism)
    end
  end
  return lmech, lcallback
end
Helper = {
  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,
  set_mechanism = function(self, mechanism)
    self.mechanism, self.callback = check_mechanism(mechanism)
    return self.mechanism
  end,
  set_callback = function(self, callback)
    if callback then
      self.callback = callback
    end
  end,
  reset_callback = function(self)
    self.callback = MECHANISMS[self.mechanism]
  end,
  reset = function(self)
    self:set_mechanism()
  end,
  get_mechanism = function(self)
    return self.mechanism
  end,
  encode = function(self, ...)
    return self.callback(...)
  end,
}
local unittest = require "unittest"
if not unittest.testing() then
  return _ENV
end
test_suite = unittest.TestSuite:new()
if HAVE_SSL then
  local _ = "ignored"
  local object = DigestMD5:new('Digest realm="test", domain="/HTTP/Digest",\z
    nonce="c8563a5b367e66b3693fbb07a53a30ba"',
    _, _, _, _)
  test_suite:add_test(unittest.keys_equal(
      object.challnvs,
      {
        nonce='c8563a5b367e66b3693fbb07a53a30ba',
        realm='test',
        domain='/HTTP/Digest',
      }
    ))
  object = DigestMD5:new('Digest nonce="9e4ab724d272474ab13b64d75300a47b", \z
    opaque="de40b82666bd5fe631a64f3b2d5a019e", \z
    realm="me@kennethreitz.com", qop=auth',
    _, _, _, _)
  test_suite:add_test(unittest.keys_equal(
      object.challnvs,
      {
        nonce='9e4ab724d272474ab13b64d75300a47b',
        opaque='de40b82666bd5fe631a64f3b2d5a019e',
        realm='me@kennethreitz.com',
        qop='auth',
      }
    ))
  object = DigestMD5:new('realm=test, domain="/HTTP/Digest",\tnonce=c8563a5b367e66b3693fbb07a53a30ba',
    _, _, _, _)
  test_suite:add_test(unittest.keys_equal(
      object.challnvs,
      {
        nonce='c8563a5b367e66b3693fbb07a53a30ba',
        realm='test',
        domain='/HTTP/Digest',
      }
    ))
end
return _ENV;
