local dns = require "dns"
local math = require "math"
local nmap = require "nmap"
local stdnse = require "stdnse"
local string = require "string"
local table = require "table"
_ENV = stdnse.module("netbios", stdnse.seeall)
types = {
  NB = 32,
  NBSTAT = 33,
}
function name_encode(name, scope)
  stdnse.debug3("Encoding name '%s'", name)
  if(#name >= 16) then
    name = string.sub(name, 1, 16)
  else
    local padding = " "
    if name == "*" then
      padding = "\0"
    end
    name = name .. string.rep(padding, 16 - #name)
  end
  name = string.upper(name)
  local L1_encoded = {}
  for i=1, #name, 1 do
    local b = string.byte(name, i)
    L1_encoded[i*2-1] = string.char(((b & 0xF0) >> 4) + 0x41)
    L1_encoded[i*2]   = string.char((b & 0x0F) + 0x41)
  end
  local L2_encoded = { string.char(32), table.concat(L1_encoded) }
  if scope ~= nil then
    local piece
    for piece in string.gmatch(scope, "[^.]+") do
      L2_encoded[#L2_encoded+1] = string.char(#piece) .. piece
    end
  end
  L2_encoded = table.concat(L2_encoded)
  stdnse.debug3("=> '%s'", L2_encoded)
  return L2_encoded
end
function name_decode(encoded_name)
  local name = ""
  local scope = ""
  local len = string.byte(encoded_name, 1)
  local i
  stdnse.debug3("Decoding name '%s'", encoded_name)
  name = name:gsub("(.)(.)", function (a, b)
      local ch = ((string.byte(a) - 0x41) << 4) | (string.byte(b) - 0x41)
      return string.char(ch)
    end)
  local pos = 34
  while #encoded_name > pos do
    local len = string.byte(encoded_name, pos)
    scope = scope .. string.sub(encoded_name, pos + 1, pos + len) .. "."
    pos = pos + 1 + len
  end
  if(#scope > 0) then
    scope = string.sub(scope, 1, #scope - 1)
  end
  stdnse.debug3("=> '%s'", name)
  return name, scope
end
function get_names(host, prefix)
  local status, names, statistics = do_nbstat(host)
  if(prefix == nil) then
    prefix = ""
  end
  if(status) then
    local result = ""
    for i = 1, #names, 1 do
      result = result .. string.format("%s%s<%02x>\n", prefix, names[i]['name'], names[i]['prefix'])
    end
    return true, result
  else
    return false, names
  end
end
function get_server_name(host, names)
  local status
  local i
  if names == nil then
    status, names = do_nbstat(host)
    if(status == false) then
      return false, names
    end
  end
  for i = 1, #names, 1 do
    if names[i]['suffix'] == 0x20 then
      return true, names[i]['name']
    end
  end
  return true, nil
end
function get_workstation_name(host, names)
  local status
  local i
  if names == nil then
    status, names = do_nbstat(host)
    if(status == false) then
      return false, names
    end
  end
  for i = 1, #names, 1 do
    if names[i]['suffix'] == 0x00 and (names[i]['flags'] & 0x8000) == 0 then
      return true, names[i]['name']
    end
  end
  return true, nil
end
function get_user_name(host, names)
  local status, server_name = get_server_name(host, names)
  if(status == false) then
    return false, server_name
  end
  if(names == nil) then
    status, names = do_nbstat(host)
    if(status == false) then
      return false, names
    end
  end
  for i = 1, #names, 1 do
    if names[i]['suffix'] == 0x03 and names[i]['name'] ~= server_name then
      return true, names[i]['name']
    end
  end
  return true, nil
end
function do_nbstat(host)
  local status, err
  local socket = nmap.new_socket()
  local encoded_name = name_encode("*")
  local statistics
  local reg
  if type(host) == "string" then --ip
    stdnse.debug3("Performing nbstat on host '%s'", host)
    nmap.registry.netbios = nmap.registry.netbios or {}
    nmap.registry.netbios[host] = nmap.registry.netbios[host] or {}
    reg = nmap.registry.netbios[host]
  else
    stdnse.debug3("Performing nbstat on host '%s'", host.ip)
    if host.registry.netbios == nil and
      nmap.registry.netbios ~= nil and
      nmap.registry.netbios[host.ip] ~= nil then
      host.registry.netbios = nmap.registry.netbios[host.ip]
    end
    host.registry.netbios = host.registry.netbios or {}
    reg = host.registry.netbios
  end
  if(reg["nbstat_names"] ~= nil) then
    stdnse.debug3(" |_ [using cached value]")
    return true, reg["nbstat_names"], reg["nbstat_statistics"]
  end
  local query = string.pack(">I2I2I2I2I2I2",
  0x1337,  -- Transaction id
  0x0000,  -- Flags
  1,       -- Questions
  0,       -- Answers
  0,       -- Authority
  0        -- Extra
  ) .. string.pack(">zI2I2",
  encoded_name, -- Encoded name
  0x0021,       -- Query type (0x21 = NBSTAT)
  0x0001        -- Class = IN
  )
  status, err = socket:connect(host, 137, "udp")
  if(status == false) then
    return false, err
  end
  status, err = socket:send(query)
  if(status == false) then
    return false, err
  end
  socket:set_timeout(1000)
  local status, result = socket:receive_bytes(1)
  if(status == false) then
    return false, result
  end
  local close_status, err = socket:close()
  if(close_status == false) then
    return false, err
  end
  if(status) then
    local pos, TRN_ID, FLAGS, QDCOUNT, ANCOUNT, NSCOUNT, ARCOUNT, rr_name, rr_type, rr_class, rr_ttl
    local rrlength, name_count
    TRN_ID, FLAGS, QDCOUNT, ANCOUNT, NSCOUNT, ARCOUNT, pos = string.unpack(">I2I2I2I2I2I2", result)
    if(TRN_ID ~= 0x1337) then
      return false, string.format("Invalid transaction ID returned: 0x%04x", TRN_ID)
    end
    if(ANCOUNT ~= 1) then
      return false, "Server returned an invalid number of answers"
    end
    if FLAGS & 0x8000 == 0 then
      return false, "Server's flags didn't indicate a response"
    end
    if FLAGS & 0x0007 ~= 0 then
      return false, string.format("Server returned a NetBIOS error: 0x%02x", FLAGS & 0x0007)
    end
    rr_name, rr_type, rr_class, rr_ttl, pos = string.unpack(">zI2I2I4", result, pos)
    if(rr_name ~= encoded_name) then
      return false, "Server returned incorrect name"
    end
    if(rr_class ~= 0x0001) then
      return false, "Server returned incorrect class"
    end
    if(rr_type ~= 0x0021) then
      return false, "Server returned incorrect query type"
    end
    rrlength, name_count, pos = string.unpack(">I2B", result, pos)
    local names = {}
    for i = 1, name_count do
      local name, suffix, flags
      name, suffix, flags, pos = string.unpack(">c15BI2", result, pos)
      name = string.gsub(name, "[ ]*$", "")
      names[i] = {}
      names[i]['name']   = name
      names[i]['suffix'] = suffix
      names[i]['flags']  = flags
      rrlength = rrlength - 18
    end
    if(rrlength > 0) then
      rrlength = rrlength - 1
    end
    statistics, pos = string.unpack(string.format(">c%d", rrlength), result, pos)
    reg["nbstat_names"] = names
    reg["nbstat_statistics"] = statistics
    return true, names, statistics
  else
    return false, "Name query failed: " .. result
  end
end
function nbquery(host, nbname, options)
  local options = options or {}
  options.port = options.port or 137
  options.retPkt = options.retPkt or true
  options.dtype = options.dtype or types.NB
  options.host = host.ip
  options.flags = options.flags or ( options.multiple and 0x0110 )
  options.id = math.random(0xFFFF)
  local encoded_name = name_encode(nbname):sub(2)
  local status, response = dns.query( encoded_name, options )
  if ( not(status) ) then return false, "ERROR: nbquery failed" end
  local results = {}
  if ( options.multiple and #response > 0 ) then
    for _, resp in ipairs(response) do
      assert( options.id == resp.output.id, "Received packet with invalid transaction ID" )
      if ( not(resp.output.answers) or #resp.output.answers < 1 ) then
        return false, "ERROR: Response contained no answers"
      end
      local dname = string.char(#resp.output.answers[1].dname) .. resp.output.answers[1].dname
      table.insert( results, { peer = resp.peer, name = name_decode(dname) } )
    end
    return true, results
  else
    local dname = string.char(#response.answers[1].dname) .. response.answers[1].dname
    return true, { { peer = host.ip, name = name_decode(dname) } }
  end
end
function flags_to_string(flags)
  local result = {}
  if flags & 0x8000 ~= 0 then
    result[#result+1] = "<group>"
  else
    result[#result+1] = "<unique>"
  end
  if flags & 0x1000 ~= 0 then
    result[#result+1] = "<deregister>"
  end
  if flags & 0x0800 ~= 0 then
    result[#result+1] = "<conflict>"
  end
  if flags & 0x0400 ~= 0 then
    result[#result+1] = "<active>"
  end
  if flags & 0x0200 ~= 0 then
    result[#result+1] = "<permanent>"
  end
  return table.concat(result)
end
return _ENV;
