local nmap = require "nmap"
local stdnse = require "stdnse"
local string = require "string"
local table = require "table"
local unicode = require "unicode"
local unittest = require "unittest"
_ENV = stdnse.module("smbauth", stdnse.seeall)
local have_ssl, openssl = pcall(require, "openssl")
local NTLMSSP_NEGOTIATE = 0x00000001
local NTLMSSP_CHALLENGE = 0x00000002
local NTLMSSP_AUTH      = 0x00000003
local session_key = string.rep("\0", 16)
local ACCOUNT_TYPES = {
  ANONYMOUS = 0,
  GUEST     = 1,
  USER      = 2,
  ADMIN     = 3
}
local function account_exists(host, username, domain)
  if(host.registry['smbaccounts'] == nil) then
    return false
  end
  for i, j in pairs(host.registry['smbaccounts']) do
    if(j['username'] == username and j['domain'] == domain) then
      return true
    end
  end
  return false
end
function next_account(host, num)
  if(num == nil) then
    if(host.registry['smbindex'] == nil) then
      host.registry['smbindex'] = 1
    else
      host.registry['smbindex'] = host.registry['smbindex'] + 1
    end
  else
    host.registry['smbindex'] = num
  end
end
function add_account(host, username, domain, password, password_hash, hash_type, is_admin)
  if(account_exists(host, username, domain)) then
    return
  end
  if(host.registry['smbaccounts'] == nil) then
    host.registry['smbaccounts'] = {}
  end
  local account_type = nil
  if(is_admin) then
    account_type = ACCOUNT_TYPES.ADMIN
  else
    if(username == '') then
      account_type = ACCOUNT_TYPES.ANONYMOUS
    elseif(string.lower(username) == 'guest') then
      account_type = ACCOUNT_TYPES.GUEST
    else
      account_type = ACCOUNT_TYPES.USER
    end
  end
  if(hash_type == nil) then
    hash_type = 'ntlm'
  end
  local new_entry = {}
  new_entry['username']      = username
  new_entry['domain']        = domain
  new_entry['password']      = password
  new_entry['password_hash'] = password_hash
  new_entry['hash_type']     = string.lower(hash_type)
  new_entry['account_type']  = account_type
  table.insert(host.registry['smbaccounts'], new_entry)
  table.sort(host.registry['smbaccounts'], function(a,b) return a['account_type'] > b['account_type'] end)
  stdnse.debug1("SMB: Added account '%s' to account list", username)
  next_account(host, 1)
end
function get_account(host)
  if(host.registry['smbindex'] == nil) then
    host.registry['smbindex'] = 1
  end
  local index = host.registry['smbindex']
  local account = host.registry['smbaccounts'][index]
  if(account == nil) then
    return false, "No accounts left to try"
  end
  return true, account['username'], account['domain'], account['password'], account['password_hash'], account['hash_type']
end
function init_account(host)
  if(host.registry['smbaccounts'] ~= nil) then
    return
  end
  host.registry['smbaccounts'] = {}
  add_account(host, '',      '', '', nil, 'none')
  if(not stdnse.get_script_args( "smbnoguest" )) then
    add_account(host, 'guest', '', '', nil, 'ntlm')
  end
  local args = nmap.registry.args
  local username      = nil
  local domain        = ''
  local password      = nil
  local password_hash = nil
  local hash_type     = 'ntlm'
  if(args.smbusername ~= nil) then
    username = args.smbusername
  elseif(args.smbuser ~= nil) then
    username = args.smbuser
  end
  if(username ~= nil) then
    if(args.smbdomain ~= nil) then
      domain = args.smbdomain
    end
    if(args.smbtype ~= nil) then
      hash_type = args.smbtype
    end
    if(args.smbpassword ~= nil) then
      password = args.smbpassword
    elseif(args.smbpass ~= nil) then
      password = args.smbpass
    end
    if(password == nil) then
      password_hash = args.smbhash
    end
    if(password == nil and password_hash == nil) then
      stdnse.debug1("SMB: Either smbpass, smbpassword, or smbhash have to be passed as script arguments to use an account")
    else
      add_account(host, username, domain, password, password_hash, hash_type)
    end
  end
end
local function lm_create_hash(password)
  if(have_ssl ~= true) then
    return false, "SMB: OpenSSL not present"
  end
  local str1, str2
  local key1, key2
  local result
  password = string.upper(password)
  do
    local buf = {}
    for i, cp in ipairs(unicode.decode(password, unicode.utf8_dec)) do
      local ch = unicode.cp437_enc(cp)
      if ch == nil then
        return false, "Couldn't encode password in CP437"
      end
      buf[i] = ch
    end
    password = table.concat(buf)
  end
  password = password .. string.rep('\0', 14 - #password)
  str1 = string.sub(password, 1, 7)
  str2 = string.sub(password, 8, 14)
  key1 = openssl.DES_string_to_key(str1)
  key2 = openssl.DES_string_to_key(str2)
  result = openssl.encrypt("DES", key1, nil, "KGS!@#$%") .. openssl.encrypt("DES", key2, nil, "KGS!@#$%")
  return true, result
end
function ntlm_create_hash(password)
  if(have_ssl ~= true) then
    return false, "SMB: OpenSSL not present"
  end
  return true, openssl.md4(unicode.utf8to16(password))
end
function lm_create_response(lanman, challenge)
  if(have_ssl ~= true) then
    return false, "SMB: OpenSSL not present"
  end
  local str1, str2, str3
  local key1, key2, key3
  local result
  lanman = lanman .. string.rep('\0', 21 - #lanman)
  str1 = string.sub(lanman, 1,  7)
  str2 = string.sub(lanman, 8,  14)
  str3 = string.sub(lanman, 15, 21)
  key1 = openssl.DES_string_to_key(str1)
  key2 = openssl.DES_string_to_key(str2)
  key3 = openssl.DES_string_to_key(str3)
  if(challenge == "") then
    stdnse.debug1("SMB: ERROR: Server returned invalid (blank) challenge value (should be 8 bytes); failing login to avoid OpenSSL crash.")
    challenge = "AAAAAAAA"
  end
  result = openssl.encrypt("DES", key1, nil, challenge) .. openssl.encrypt("DES", key2, nil, challenge) .. openssl.encrypt("DES", key3, nil, challenge)
  return true, result
end
function ntlm_create_response(ntlm, challenge)
  if(have_ssl ~= true) then
    return false, "SMB: OpenSSL not present"
  end
  return lm_create_response(ntlm, challenge)
end
function ntlm_create_mac_key(ntlm_hash, ntlm_response, is_extended)
  if(have_ssl ~= true) then
    return false, "SMB: OpenSSL not present"
  end
  if(is_extended) then
    return openssl.md4(ntlm_hash)
  else
    return openssl.md4(ntlm_hash) .. ntlm_response
  end
end
function lm_create_mac_key(lm_hash, lm_response, is_extended)
  if(have_ssl ~= true) then
    return false, "SMB: OpenSSL not present"
  end
  if(is_extended) then
    return string.sub(lm_hash, 1, 8) .. string.rep('\0', 8)
  else
    return string.sub(lm_hash, 1, 8) .. string.rep('\0', 8) .. lm_response
  end
end
function ntlmv2_create_hash(ntlm, username, domain)
  if(have_ssl ~= true) then
    return false, "SMB: OpenSSL not present"
  end
  username = unicode.utf8to16(string.upper(username))
  domain   = unicode.utf8to16(string.upper(domain))
  return true, openssl.hmac("MD5", ntlm, username .. domain)
end
function lmv2_create_response(ntlm, username, domain, challenge)
  if(have_ssl ~= true) then
    return false, "SMB: OpenSSL not present"
  end
  return ntlmv2_create_response(ntlm, username, domain, challenge, 8)
end
function ntlmv2_create_response(ntlm, username, domain, challenge, client_challenge_length)
  if(have_ssl ~= true) then
    return false, "SMB: OpenSSL not present"
  end
  local client_challenge = openssl.rand_bytes(client_challenge_length)
  local status, ntlmv2_hash = ntlmv2_create_hash(ntlm, username, domain)
  return true, openssl.hmac("MD5", ntlmv2_hash, challenge .. client_challenge) .. client_challenge
end
function ntlmv2_session_response(ntlm_password_hash, challenge)
  local client_nonce = openssl.rand_bytes(8)
  local lm_response = client_nonce .. string.rep('\0', 24 - #client_nonce)
  local session_nonce = challenge .. client_nonce
  local ntlm_session_hash  = openssl.md5(session_nonce):sub(1,8)
  local status, ntlm_response =  lm_create_response(ntlm_password_hash, ntlm_session_hash)
  return status, lm_response, ntlm_response
end
function get_password_response(ip, username, domain, password, password_hash, hash_type, challenge, is_extended)
  local status
  local lm_hash   = nil
  local ntlm_hash = nil
  local mac_key   = nil
  local lm_response, ntlm_response
  if(password == nil and password_hash == nil) then
    stdnse.debug2("SMB: Couldn't find password or hash to use (assuming blank)")
    password = ""
  end
  if(hash_type == 'none') then
    return '\0', '', nil
  end
  if(password ~= nil) then
    status, lm_hash   = lm_create_hash(password)
    status, ntlm_hash = ntlm_create_hash(password)
  else
    if(password_hash ~= nil) then
      if(string.find(password_hash, "^" .. string.rep("%x%x", 16) .. "$")) then
        stdnse.debug2("SMB: Found a 16-byte hex string")
        lm_hash   = stdnse.fromhex(password_hash:sub(1, 32))
        ntlm_hash = stdnse.fromhex(password_hash:sub(1, 32))
      elseif(string.find(password_hash, "^" .. string.rep("%x%x", 32) .. "$")) then
        stdnse.debug2("SMB: Found a 32-byte hex string")
        lm_hash   = stdnse.fromhex(password_hash:sub(1, 32))
        ntlm_hash = stdnse.fromhex(password_hash:sub(33, 64))
      elseif(string.find(password_hash, "^" .. string.rep("%x%x", 16) .. "." .. string.rep("%x%x", 16) .. "$")) then
        stdnse.debug2("SMB: Found two 16-byte hex strings")
        lm_hash   = stdnse.fromhex(password_hash:sub(1, 32))
        ntlm_hash = stdnse.fromhex(password_hash:sub(34, 65))
      else
        stdnse.debug1("SMB: ERROR: Hash(es) provided in an invalid format (should be 32, 64, or 65 hex characters)")
        lm_hash = nil
        ntlm_hash = nil
      end
    end
  end
  if(lm_hash == nil or ntlm_hash == nil) then
    stdnse.debug2("SMB: Couldn't determine which password to use, using a blank one")
    return "", ""
  end
  stdnse.debug2("SMB: Lanman hash: %s", stdnse.tohex(lm_hash))
  stdnse.debug2("SMB: NTLM   hash: %s", stdnse.tohex(ntlm_hash))
  if(hash_type == "v1") then
    stdnse.debug2("SMB: Creating v1 response")
    status, lm_response   = lm_create_response(lm_hash, challenge)
    status, ntlm_response = ntlm_create_response(ntlm_hash, challenge)
    mac_key               = ntlm_create_mac_key(ntlm_hash, ntlm_response, is_extended)
  elseif(hash_type == "lm") then
    stdnse.debug2("SMB: Creating LMv1 response")
    status, lm_response   = lm_create_response(lm_hash, challenge)
    ntlm_response = ""
    mac_key               = lm_create_mac_key(lm_hash, lm_response, is_extended)
  elseif(hash_type == "ntlm") then
    stdnse.debug2("SMB: Creating NTLMv1 response")
    status, lm_response   = ntlm_create_response(ntlm_hash, challenge)
    status, ntlm_response = ntlm_create_response(ntlm_hash, challenge)
    mac_key               = ntlm_create_mac_key(ntlm_hash, ntlm_response, is_extended)
  elseif(hash_type == "v2") then
    stdnse.debug2("SMB: Creating v2 response")
    status, lm_response   = lmv2_create_response(ntlm_hash, username, domain, challenge)
    status, ntlm_response = ntlmv2_create_response(ntlm_hash, username, domain, challenge, 24)
  elseif(hash_type == "lmv2") then
    stdnse.debug2("SMB: Creating LMv2 response")
    status, lm_response   = lmv2_create_response(ntlm_hash, username, domain, challenge)
    ntlm_response = ""
  elseif(hash_type == "ntlmv2_session") then
    stdnse.debug2("SMB: Creating nltmv2 session response")
    status, lm_response, ntlm_response = ntlmv2_session_response(ntlm_hash, challenge)
  else
    if(hash_type ~= nil) then
      stdnse.debug1("SMB: Invalid login type specified ('%s'), using default (NTLM)", hash_type)
    else
      stdnse.debug1("SMB: No login type specified, using default (NTLM)")
    end
    status, lm_response   = ntlm_create_response(ntlm_hash, challenge)
    status, ntlm_response = ntlm_create_response(ntlm_hash, challenge)
  end
  stdnse.debug2("SMB: Lanman response: %s", stdnse.tohex(lm_response))
  stdnse.debug2("SMB: NTLM   response: %s", stdnse.tohex(ntlm_response))
  return lm_response, ntlm_response, mac_key
end
function get_security_blob(security_blob, ip, username, domain, password, password_hash, hash_type, flags)
  local pos = 1
  local new_blob
  local flags = flags or 0x00008215 -- (NEGOTIATE_SIGN_ALWAYS | NEGOTIATE_NTLM | NEGOTIATE_SIGN | REQUEST_TARGET | NEGOTIATE_UNICODE)
  if(security_blob == nil) then
    new_blob = string.pack("<zI4I4I8I8",
    "NTLMSSP",            -- Identifier
    NTLMSSP_NEGOTIATE,    -- Type
    flags,                -- Flags
    0,                    -- Calling workstation domain
    0                     -- Calling workstation name
    )
    return true, new_blob, "", ""
  else
    local identifier, message_type, domain_length, domain_max, domain_offset, server_flags, challenge, reserved = string.unpack("<I8I4I2I2I4I4c8c8", security_blob)
    local lanman, ntlm, mac_key = get_password_response(ip, username, domain, password, password_hash, hash_type, challenge, true)
    local hostname = unicode.utf8to16("nmap")
    username = unicode.utf8to16(username)
    domain   = (#username > 0 ) and unicode.utf8to16(domain) or ""
    ntlm     = (#username > 0 ) and ntlm or ""
    lanman   = (#username > 0 ) and lanman or '\0'
    local domain_offset = 0x40
    local username_offset = domain_offset + #domain
    local hostname_offset = username_offset + #username
    local lanman_offset = hostname_offset + #hostname
    local ntlm_offset = lanman_offset + #lanman
    local sessionkey_offset = ntlm_offset + #ntlm
    new_blob = string.pack("<zI4 I2I2I4 I2I2I4 I2I2I4 I2I2I4 I2I2I4 I2I2I4 I4",
      "NTLMSSP",
      NTLMSSP_AUTH,
      #lanman,
      #lanman,
      lanman_offset,
      ( #ntlm > 0 and #ntlm - 16 or 0 ),
      ( #ntlm > 0 and #ntlm - 16 or 0 ),
      ntlm_offset,
      #domain,
      #domain,
      domain_offset,
      #username,
      #username,
      username_offset,
      #hostname,
      #hostname,
      hostname_offset,
      #session_key,
      #session_key,
      sessionkey_offset,
      flags)
      .. domain
      .. username
      .. hostname
      .. lanman
      .. ntlm
      .. session_key
    return true, new_blob, mac_key
  end
end
function get_host_info_from_security_blob(security_blob)
  local identifier, message_type, domain_length, domain_max, domain_offset, server_flags, challenge, hpos = string.unpack("<c8I4 I2I2I4 I4I8", security_blob)
  if ( identifier ~= "NTLMSSP\0" ) then
    stdnse.debug1("SMB: Invalid NTLM challenge message: unexpected signature." )
    return false, "Invalid NTLM challenge message"
  elseif ( message_type ~= 0x2 ) then
    stdnse.debug1("SMB: Invalid NTLM challenge message: unexpected message type: %d.", message_type )
    return false, "Invalid message type in NTLM challenge message"
  end
  local ntlm_challenge = {}
  if ( domain_length > 0 ) then
    local length = domain_length
    local pos = domain_offset + 1 -- +1 to convert to Lua's 1-based indexes
    local target_realm
    target_realm = string.unpack("c" .. length, security_blob, pos )
    ntlm_challenge[ "target_realm" ] = unicode.utf16to8( target_realm )
  end
  if hpos + domain_length > #security_blob then
    return ntlm_challenge
  end
  local context, target_info_length, target_info_max, target_info_offset, hpos = string.unpack("<I8 I2I2I4", security_blob, hpos)
  if target_info_offset >= hpos + 7 and domain_offset >= hpos + 7 then
    local major, minor, build, reserved = string.unpack("<BBI2c4", security_blob, hpos)
    if reserved == "\0\0\0\x0f" then
      ntlm_challenge.os_major_version = major
      ntlm_challenge.os_minor_version = minor
      ntlm_challenge.os_build = build
    else
      stdnse.debug2("smbauth: Unknown OS info structure in NTLM handshake")
    end
  end
  if ( target_info_length > 0 ) then
    local NTLM_AV_ID_VALUES = {
      MsvAvEOL = 0x0,
      MsvAvNbComputerName = 0x1,
      MsvAvNbDomainName = 0x2,
      MsvAvDnsComputerName = 0x3,
      MsvAvDnsDomainName = 0x4,
      MsvAvDnsTreeName = 0x5,
      MsvAvFlags = 0x6,
      MsvAvTimestamp = 0x7,
      MsvAvRestrictions = 0x8,
      MsvAvTargetName = 0x9,
      MsvAvChannelBindings = 0xA,
    }
    local NTLM_AV_ID_NAMES = {
      [NTLM_AV_ID_VALUES.MsvAvNbComputerName] = "netbios_computer_name",
      [NTLM_AV_ID_VALUES.MsvAvNbDomainName] = "netbios_domain_name",
      [NTLM_AV_ID_VALUES.MsvAvDnsComputerName] = "fqdn",
      [NTLM_AV_ID_VALUES.MsvAvDnsDomainName] = "dns_domain_name",
      [NTLM_AV_ID_VALUES.MsvAvDnsTreeName] = "dns_forest_name",
      [NTLM_AV_ID_VALUES.MsvAvTimestamp] = "timestamp",
    }
    local length = target_info_length
    local pos = target_info_offset + 1 -- +1 to convert to Lua's 1-based indexes
    local target_info
    target_info = string.unpack("c" .. length, security_blob, pos)
    pos = 1 -- reset pos to 1, since we'll be working out of just the target_info
    repeat
      local value, av_id
      av_id, value, pos = string.unpack( "<I2s2", target_info, pos )
      local friendly_name = NTLM_AV_ID_NAMES[ av_id ]
      if ( av_id == NTLM_AV_ID_VALUES.MsvAvEOL ) then
        break
      elseif ( av_id == NTLM_AV_ID_VALUES.MsvAvTimestamp ) then
        ntlm_challenge[ friendly_name ] = string.unpack( "<I8", value )
      elseif ( friendly_name ) then
        ntlm_challenge[ friendly_name ] = unicode.utf16to8( value )
      end
    until ( pos >= #target_info )
  end
  return ntlm_challenge
end
function calculate_signature(mac_key, data)
  if(have_ssl) then
    return string.sub(openssl.md5(mac_key .. data), 1, 8)
  else
    return string.rep('\0', 8)
  end
end
if not unittest.testing() then
  return _ENV
end
test_suite = unittest.TestSuite:new()
if have_ssl then
  test_suite:add_test(unittest.equal(
      stdnse.tohex(select(-1, lm_create_hash("passphrase"))),
      "855c3697d9979e78ac404c4ba2c66533"
      ),
    "lm_create_hash"
    )
  test_suite:add_test(unittest.equal(
      stdnse.tohex(select(-1, ntlm_create_hash("passphrase"))),
      "7f8fe03093cc84b267b109625f6bbf4b"
      ),
    "ntlm_create_hash"
    )
  test_suite:add_test(unittest.equal(
      stdnse.tohex(select(-1, lm_create_hash("ÅÇÅÇ"))),
      "1830f5732b438091aad3b435b51404ee"
      ),
    "lm_create_hash"
    )
  test_suite:add_test(unittest.equal(
      stdnse.tohex(select(-1, ntlm_create_hash("öäü"))),
      "4848bcb81cf018c3b70ea1479bd1374d"
      ),
    "ntlm_create_hash"
    )
end
return _ENV;
