local asn1 = require "asn1"
local datetime = require "datetime"
local io = require "io"
local nmap = require "nmap"
local os = require "os"
local stdnse = require "stdnse"
local string = require "string"
local stringaux = require "stringaux"
local table = require "table"
local comm = require "comm"
_ENV = stdnse.module("ldap", stdnse.seeall)
local ldapMessageId = 1
ERROR_MSG = {}
ERROR_MSG[1]  = "Initialization of LDAP library failed."
ERROR_MSG[4]  = "Size limit exceeded."
ERROR_MSG[13] = "Confidentiality required"
ERROR_MSG[32] = "No such object"
ERROR_MSG[34] = "Invalid DN"
ERROR_MSG[49] = "The supplied credential is invalid."
ERRORS = {
  LDAP_SUCCESS = 0,
  LDAP_SIZELIMIT_EXCEEDED = 4
}
APPNO = {
  BindRequest = 0,
  BindResponse = 1,
  UnbindRequest = 2,
  SearchRequest = 3,
  SearchResponse = 4,
  SearchResDone = 5
}
FILTER = {
  _and = 0,
  _or = 1,
  _not = 2,
  equalityMatch = 3,
  substrings = 4,
  greaterOrEqual = 5,
  lessOrEqual = 6,
  present = 7,
  approxMatch = 8,
  extensibleMatch = 9
}
SCOPE = {
  base=0,
  one=1,
  sub= 2,
  children=3,
  default = 0
}
DEREFPOLICY = {
  never=0,
  searching=1,
  finding = 2,
  always=3,
  default = 0
}
local tagEncoder = {}
tagEncoder['table'] = function(self, val)
  if (val._ldap == '\x0A') then
    local ival = self.encodeInt(val[1])
    local len = self.encodeLength(#ival)
    return val._ldap .. len .. ival
  end
  if (val._ldaptype) then
    local len
    if val[1] == nil or #val[1] == 0 then
      return val._ldaptype .. '\0'
    else
      len = self.encodeLength(#val[1])
      return val._ldaptype .. len .. val[1]
    end
  end
  local encVal = ""
  for _, v in ipairs(val) do
    encVal = encVal .. encode(v) -- todo: buffer?
  end
  local tableType = val._snmp or "\x30"
  return tableType .. self.encodeLength(#encVal) .. encVal
end
function encode(val)
  local encoder = asn1.ASN1Encoder:new()
  local encValue
  encoder:registerTagEncoders(tagEncoder)
  encValue = encoder:encode(val)
  if encValue then
    return encValue
  end
  return ''
end
local tagDecoder = {}
tagDecoder["\x0A"] = function( self, encStr, elen, pos )
  return self.decodeInt(encStr, elen, pos)
end
tagDecoder["\x8A"] = function( self, encStr, elen, pos )
  return string.unpack("c" .. elen, encStr, pos)
end
tagDecoder["\x31"] = function( self, encStr, elen, pos )
  return nil, pos
end
function decode(encStr, pos)
  local decoder = asn1.ASN1Decoder:new()
  decoder:registerTagDecoders( tagDecoder )
  return decoder:decode( encStr, pos )
end
local function decodeSeq(encStr, len, pos)
  local seq = {}
  local sPos = 1
  if #encStr - pos + 1 < len then
    return seq, nil
  end
  local sStr, newpos = string.unpack("c" .. len, encStr, pos)
  while (sPos < len) do
    local newSeq
    newSeq, sPos = decode(sStr, sPos)
    table.insert(seq, newSeq)
  end
  return seq, newpos
end
function encodeLDAPOp( appno, isConstructed, data )
  local encoded_str = ""
  local asn1_type = asn1.BERtoInt( asn1.BERCLASS.Application, isConstructed, appno )
  encoded_str = encode( { _ldaptype = string.pack("B", asn1_type), data } )
  return encoded_str
end
function searchRequest( socket, params )
  local searchResEntries = { errorMessage="", resultCode = 0}
  local catch = function() socket:close() stdnse.debug1("SearchRequest failed") end
  local try = nmap.new_try(catch)
  local attributes = params.attributes
  local request = encode(params.baseObject)
  local attrSeq = ''
  local requestData, messageSeq, data
  local maxObjects = params.maxObjects or -1
  local encoder = asn1.ASN1Encoder:new()
  local decoder = asn1.ASN1Decoder:new()
  encoder:registerTagEncoders(tagEncoder)
  decoder:registerTagDecoders(tagDecoder)
  request = request .. encode( { _ldap='\x0A', params.scope } )--scope
  request = request .. encode( { _ldap='\x0A', params.derefPolicy } )--derefpolicy
  request = request .. encode( params.sizeLimit or 0)--sizelimit
  request = request .. encode( params.timeLimit or 0)--timelimit
  request = request .. encode( params.typesOnly or false)--TypesOnly
  if params.filter then
    request = request .. createFilter( params.filter )
  else
    request = request .. encode( { _ldaptype='\x87', "objectclass" } )-- filter : string, presence
  end
  if  attributes~= nil then
    for _,attr in ipairs(attributes) do
      attrSeq = attrSeq .. encode(attr)
    end
  end
  request = request .. encoder:encodeSeq(attrSeq)
  requestData = encodeLDAPOp(APPNO.SearchRequest, true, request)
  messageSeq = encode(ldapMessageId)
  ldapMessageId = ldapMessageId +1
  messageSeq = messageSeq .. requestData
  data = encoder:encodeSeq(messageSeq)
  try( socket:send( data ) )
  data = ""
  while true do
    local len, pos, messageId = 0, 2, -1
    local tmp = ""
    local _, objectName, attributes, ldapOp
    local attributes
    local searchResEntry = {}
    if ( maxObjects == 0 ) then
      break
    elseif ( maxObjects > 0 ) then
      maxObjects = maxObjects - 1
    end
    if data:len() > 6 then
      len, pos = decoder.decodeLength( data, pos )
    else
      data = data .. try( socket:receive() )
      len, pos = decoder.decodeLength( data, pos )
    end
    while ( len + pos - 1 > data:len() ) do
      data = data .. try( socket:receive() )
    end
    messageId, pos = decode( data, pos )
    tmp, pos = string.unpack("B", data, pos)
    len, pos = decoder.decodeLength( data, pos )
    ldapOp = asn1.intToBER( tmp )
    searchResEntry = {}
    if ldapOp.number == APPNO.SearchResDone then
      searchResEntry.resultCode, pos = decode( data, pos )
      if ( searchResEntry.resultCode ~= 0 ) then
        local error_msg
        searchResEntry.matchedDN, pos = decode( data, pos )
        searchResEntry.errorMessage, pos = decode( data, pos )
        error_msg = ERROR_MSG[searchResEntry.resultCode]
        if #searchResEntries == 0 then
          return false, string.format("Code: %d|Error: %s|Details: %s", searchResEntry.resultCode, error_msg or "", searchResEntry.errorMessage or "" )
        else
          searchResEntries.errorMessage = string.format("Code: %d|Error: %s|Details: %s", searchResEntry.resultCode, error_msg or "", searchResEntry.errorMessage or "" )
          searchResEntries.resultCode = searchResEntry.resultCode
          return true, searchResEntries
        end
      end
      break
    end
    searchResEntry.objectName, pos = decode( data, pos )
    if ldapOp.number == APPNO.SearchResponse then
      searchResEntry.attributes, pos = decode( data, pos )
      table.insert( searchResEntries, searchResEntry )
    end
    if data:len() > pos then
      data = data:sub(pos)
    else
      data = ""
    end
  end
  return true, searchResEntries
end
function udpSearchRequest( host, port, params )
  local searchResEntries = { errorMessage="", resultCode = 0}
  local catch = function() stdnse.debug1("udpSearchRequest failed") end
  local try = nmap.new_try(catch)
  local attributes = params.attributes
  local request = encode(params.baseObject)
  local attrSeq = ''
  local requestData, messageSeq, data
  local maxObjects = params.maxObjects or -1
  local encoder = asn1.ASN1Encoder:new()
  local decoder = asn1.ASN1Decoder:new()
  encoder:registerTagEncoders(tagEncoder)
  decoder:registerTagDecoders(tagDecoder)
  request = request .. encode( { _ldap='\x0A', params.scope } )--scope
  request = request .. encode( { _ldap='\x0A', params.derefPolicy } )--derefpolicy
  request = request .. encode( params.sizeLimit or 0)--sizelimit
  request = request .. encode( params.timeLimit or 0)--timelimit
  request = request .. encode( params.typesOnly or false)--TypesOnly
  if params.filter then
    request = request .. createFilter( params.filter )
  else
    request = request .. encode( { _ldaptype='\x87', "objectclass" } )-- filter : string, presence
  end
  if  attributes~= nil then
    for _,attr in ipairs(attributes) do
      attrSeq = attrSeq .. encode(attr)
    end
  end
  request = request .. encoder:encodeSeq(attrSeq)
  requestData = encodeLDAPOp(APPNO.SearchRequest, true, request)
  messageSeq = encode(ldapMessageId)
  ldapMessageId = ldapMessageId +1
  messageSeq = messageSeq .. requestData
  data = encoder:encodeSeq(messageSeq)
  local status, response = comm.exchange(host, port, data)
  while true do
    local len, pos, messageId = 0, 0, -1
    local tmp = ""
    local _, objectName, attributes, ldapOp
    local attributes
    local searchResEntry = {}
    if ( maxObjects == 0 ) then
      break
    elseif ( maxObjects > 0 ) then
      maxObjects = maxObjects - 1
    end
    tmp, pos = string.unpack("B", response, pos)
    len, pos = decoder.decodeLength( response, pos )
    messageId, pos = decode( response, pos )
    tmp, pos = string.unpack("B", response, pos)
    len, pos = decoder.decodeLength( response, pos )
    ldapOp = asn1.intToBER( tmp )
    searchResEntry = {}
    if ldapOp.number == APPNO.SearchResDone then
      searchResEntry.resultCode, pos = decode( response, pos )
      if ( searchResEntry.resultCode ~= 0 ) then
        local error_msg
        searchResEntry.matchedDN, pos = decode( response, pos )
        searchResEntry.errorMessage, pos = decode( response, pos )
        error_msg = ERROR_MSG[searchResEntry.resultCode]
        if #searchResEntries == 0 then
          return false, string.format("Code: %d|Error: %s|Details: %s", searchResEntry.resultCode, error_msg or "", searchResEntry.errorMessage or "" )
        else
          searchResEntries.errorMessage = string.format("Code: %d|Error: %s|Details: %s", searchResEntry.resultCode, error_msg or "", searchResEntry.errorMessage or "" )
          searchResEntries.resultCode = searchResEntry.resultCode
          return true, searchResEntries
        end
      end
      break
    end
    searchResEntry.objectName, pos = decode( response, pos )
    if ldapOp.number == APPNO.SearchResponse then
      searchResEntry.attributes, pos = decode( response, pos )
      table.insert( searchResEntries, searchResEntry )
    end
    if response:len() > pos then
      response = response:sub(pos)
    else
      response = ""
    end
  end
  return true, searchResEntries
end
function bindRequest( socket, params )
  local catch = function() socket:close() stdnse.debug1("bindRequest failed") end
  local try = nmap.new_try(catch)
  local ldapAuth = encode( { _ldaptype = '\x80', params.password } )
  local bindReq = encode( params.version ) .. encode( params.username ) .. ldapAuth
  local ldapMsg = encode(ldapMessageId) .. encodeLDAPOp( APPNO.BindRequest, true, bindReq )
  local packet
  local pos, packet_len, resultCode, tmp, len, _
  local response = {}
  local encoder = asn1.ASN1Encoder:new()
  local decoder = asn1.ASN1Decoder:new()
  encoder:registerTagEncoders(tagEncoder)
  decoder:registerTagDecoders(tagDecoder)
  packet = encoder:encodeSeq( ldapMsg )
  ldapMessageId = ldapMessageId +1
  try( socket:send( packet ) )
  packet = try( socket:receive() )
  packet_len, pos = decoder.decodeLength( packet, 2 )
  response.messageID, pos = decode( packet, pos )
  tmp, pos = string.unpack("B", packet, pos)
  len, pos = decoder.decodeLength( packet, pos )
  response.protocolOp = asn1.intToBER( tmp )
  if response.protocolOp.number ~= APPNO.BindResponse then
    return false, string.format("Received incorrect Op in packet: %d, expected %d", response.protocolOp.number, APPNO.BindResponse)
  end
  response.resultCode, pos = decode( packet, pos )
  if ( response.resultCode ~= 0 ) then
    local error_msg
    response.matchedDN, pos = decode( packet, pos )
    response.errorMessage, pos = decode( packet, pos )
    error_msg = ERROR_MSG[response.resultCode]
    return false, string.format("\n  Error: %s\n  Details: %s",
      error_msg or "Unknown error occurred (code: " .. response.resultCode ..
      ")", response.errorMessage or "" )
  else
    return true, "Success"
  end
end
function unbindRequest( socket )
  local ldapMsg, packet
  local catch = function() socket:close() stdnse.debug1("bindRequest failed") end
  local try = nmap.new_try(catch)
  local encoder = asn1.ASN1Encoder:new()
  encoder:registerTagEncoders(tagEncoder)
  ldapMessageId = ldapMessageId +1
  ldapMsg = encode( ldapMessageId ) .. encodeLDAPOp( APPNO.UnbindRequest, false, nil)
  packet = encoder:encodeSeq( ldapMsg )
  try( socket:send( packet ) )
  return true, ""
end
function createFilter( filter )
  local asn1_type = asn1.BERtoInt( asn1.BERCLASS.ContextSpecific, true, filter.op )
  local filter_str = ""
  if type(filter.val) == 'table' then
    for _, v in ipairs( filter.val ) do
      filter_str = filter_str .. createFilter( v )
    end
  else
    local obj = encode( filter.obj )
    local val = ''
    if ( filter.op == FILTER['substrings'] ) then
      local tmptable = stringaux.strsplit('*', filter.val)
      local tmp_result = ''
      if (#tmptable <= 1 ) then
        tmp_result = string.pack('Bs1', 0x81, filter.val)
      else
        for indexval, substr in ipairs(tmptable) do
          if (indexval == 1) and (substr ~= '') then
            tmp_result = '\x80' .. string.char(#substr) .. substr
          end
          if (indexval ~= #tmptable) and (indexval ~= 1) and (substr ~= '') then
            tmp_result = tmp_result .. string.pack('Bs1', 0x81, substr)
          end
          if (indexval == #tmptable) and (substr ~= '') then
            tmp_result = tmp_result .. string.pack('Bs1', 0x82, substr)
          end
        end
      end
      val = asn1.ASN1Encoder:encodeSeq( tmp_result )
    elseif ( filter.op == FILTER['extensibleMatch'] ) then
      local tmptable = stringaux.strsplit(':=', filter.val)
      local tmp_result = ''
      local OID, bitmask
      if ( tmptable[1] ~= nil ) then
        OID = tmptable[1]
      else
        return false, ("ERROR: Invalid extensibleMatch query format")
      end
      if ( tmptable[2] ~= nil ) then
        bitmask = tmptable[2]
      else
        return false, ("ERROR: Invalid extensibleMatch query format")
      end
      tmp_result = string.pack('Bs1 Bs1 Bs1 Bs1',
        0x81, OID,
        0x82, filter.obj,
        0x83, bitmask,
        0x84, '\x00')
      return '\xa9' .. asn1.ASN1Encoder.encodeLength(#tmp_result) .. tmp_result
    else
      val = encode( filter.val )
    end
    filter_str = filter_str .. obj .. val
  end
  return encode( { _ldaptype=string.pack("B", asn1_type), filter_str } )
end
function searchResultToTable( searchEntries )
  local result = {}
  for _, v in ipairs( searchEntries ) do
    local result_part = {}
    if v.objectName and v.objectName:len() > 0 then
      result_part.name = string.format("dn: %s", v.objectName)
    else
      result_part.name = "<ROOT>"
    end
    local attribs = {}
    if ( v.attributes ~= nil ) then
      for _, attrib in ipairs( v.attributes ) do
        for i=2, #attrib do
          if ( attrib[1] == "objectSid" ) then
            table.insert( attribs, string.format( "%s: %s", attrib[1], convertObjectSid( attrib[i] ) ) )
          elseif ( attrib[1] == "objectGUID") then
            local o = {string.unpack(("B"):rep(16), attrib[i] )}
            table.insert( attribs, string.format( "%s: %x%x%x%x-%x%x-%x%x-%x%x-%x%x%x%x%x%x",
              attrib[1], o[4], o[3], o[2], o[1], table.unpack(o, 5, 16)))
          elseif ( attrib[1] == "lastLogon" or attrib[1] == "lastLogonTimestamp" or attrib[1] == "pwdLastSet" or attrib[1] == "accountExpires" or attrib[1] == "badPasswordTime"  ) then
            table.insert( attribs, string.format( "%s: %s", attrib[1], convertADTimeStamp(attrib[i]) ) )
          elseif ( attrib[1] == "whenChanged" or attrib[1] == "whenCreated" or attrib[1] == "dSCorePropagationData" ) then
            table.insert( attribs, string.format( "%s: %s", attrib[1], convertZuluTimeStamp(attrib[i]) ) )
          else
            table.insert( attribs, string.format( "%s: %s", attrib[1], attrib[i] ) )
          end
        end
      end
      table.insert( result_part, attribs )
    end
    table.insert( result, result_part )
  end
  return result
end
function searchResultToFile( searchEntries, filename )
  local f = io.open( filename, "w")
  if ( not(f) ) then
    return false, ("ERROR: Failed to open file (%s)"):format(filename)
  end
  local attrib_table = {}
  for _, v in ipairs( searchEntries ) do
    if ( v.attributes ~= nil ) then
      for _, attrib in ipairs( v.attributes ) do
        for i=2, #attrib do
          if ( attrib_table[attrib[1]] == nil ) then
            attrib_table[attrib[1]] = ''
          end
        end
      end
    end
  end
  local host_table = {}
  for _, v in ipairs( searchEntries ) do
    if v.objectName and v.objectName:len() > 0 then
      local host = {}
      if v.objectName and v.objectName:len() > 0 then
        host_table[v.objectName] = {attributes = copyTable(attrib_table) }
      end
    end
  end
  for _, v in ipairs( searchEntries ) do
    if ( v.attributes ~= nil ) then
      for _, attrib in ipairs( v.attributes ) do
        for i=2, #attrib do
          if ( attrib[1] == "objectSid" ) then
            host_table[string.format("%s", v.objectName)].attributes[attrib[1]] = string.format( "%s", convertObjectSid(attrib[i]))
          elseif ( attrib[1] == "objectGUID") then
            local o = {string.unpack(("B"):rep(16), attrib[i] )}
            host_table[string.format("%s", v.objectName)].attributes[attrib[1]] = string.format(
              "%s: %x%x%x%x-%x%x-%x%x-%x%x-%x%x%x%x%x%x",
              attrib[1], o[4], o[3], o[2], o[1], table.unpack(o, 5, 16))
          elseif ( attrib[1] == "lastLogon" or attrib[1] == "lastLogonTimestamp" or attrib[1] == "pwdLastSet" or attrib[1] == "accountExpires" or attrib[1] == "badPasswordTime" ) then
            host_table[string.format("%s", v.objectName)].attributes[attrib[1]] = convertADTimeStamp(attrib[i])
          elseif ( attrib[1] == "whenChanged" or attrib[1] == "whenCreated" or attrib[1] == "dSCorePropagationData" ) then
            host_table[string.format("%s", v.objectName)].attributes[attrib[1]] = convertZuluTimeStamp(attrib[i])
          else
            host_table[v.objectName].attributes[attrib[1]] = string.format( "%s", attrib[i] )
          end
        end
      end
    end
  end
  local output = "\"name\""
  for attribute, value in pairs(attrib_table) do
    output = output .. ",\"" .. attribute .. "\""
  end
  output = output .. "\n"
  for name, attribs in pairs(host_table) do
    output = output .. "\"" .. name .. "\""
    local host_attribs = attribs.attributes
    for attribute, value in pairs(attrib_table) do
      output = output .. ",\"" .. host_attribs[attribute] .. "\""
    end
    output = output .. "\n"
  end
  if ( not(f:write( output .."\n" ) ) ) then
    f:close()
    return false, ("ERROR: Failed to write file (%s)"):format(filename)
  end
  f:close()
  return true
end
function extractAttribute( searchEntries, attributeName )
  local attributeTbl = {}
  for _, v in ipairs( searchEntries ) do
    if ( v.attributes ~= nil ) then
      for _, attrib in ipairs( v.attributes ) do
        local attribType = attrib[1]
        for i=2, #attrib do
          if ( attribType:upper() == attributeName:upper() ) then
            table.insert( attributeTbl, attrib[i])
          end
        end
      end
    end
  end
  return ( #attributeTbl > 0 and attributeTbl or nil )
end
function convertADTimeStamp(timestamp)
  local result = 0
  local base_time = tonumber(os.time({year=1601, month=1, day=1, hour=0, minute=0, sec =0}))
  timestamp = tonumber(timestamp)
  if (timestamp and timestamp > 0) then
    result = ( timestamp //  10000000 ) - 3036
    result = result + base_time
    result = datetime.format_timestamp(result, 0)
  else
    result = 'Never'
  end
  return result
end
function convertZuluTimeStamp(timestamp)
  if ( type(timestamp) == 'string' and string.sub(timestamp,-3) == '.0Z' ) then
    local year  = string.sub(timestamp,1,4)
    local month = string.sub(timestamp,5,6)
    local day   = string.sub(timestamp,7,8)
    local hour  = string.sub(timestamp,9,10)
    local mins  = string.sub(timestamp,11,12)
    local secs  = string.sub(timestamp,13,14)
    local result = year .. "/" .. month .. "/" .. day .. " " .. hour .. ":" .. mins .. ":" .. secs .. " UTC"
    return result
  else
    return 'Invalid date format'
  end
end
function convertObjectSid(data)
  local pos, revision, auth, sub_auth_size, sub_auth, result
  revision, pos = string.unpack('I1', data, 1)
  sub_auth_size, pos = string.unpack('I1', data, pos)
  auth, pos = string.unpack('>I6', data, pos)
  sub_auth = ''
  local tmp
  local cnt = 0
  while (cnt < sub_auth_size) do
    tmp, pos = string.unpack('<I4', data, pos)
    sub_auth = sub_auth .. '-' .. tmp
    cnt = cnt + 1
  end
  result = revision .. '-' .. auth .. sub_auth
  return result
end
function copyTable(targetTable)
  local temp = { }
  for key, val in pairs(targetTable) do
    temp[key] = val
  end
  return setmetatable(temp, getmetatable(targetTable))
end
return _ENV;
