local coroutine = require "coroutine"
local io = require "io"
local ipOps = require "ipOps"
local nmap = require "nmap"
local stdnse = require "stdnse"
local stringaux = require "stringaux"
local table = require "table"
local tableaux = require "tableaux"
_ENV = stdnse.module("creds", stdnse.seeall)
State = {
  LOCKED = 1,
  VALID = 2,
  DISABLED = 4,
  CHANGEPW = 8,
  PARAM = 16,
  EXPIRED = 32,
  TIME_RESTRICTED = 64,
  HOST_RESTRICTED = 128,
  LOCKED_VALID = 256,
  DISABLED_VALID = 512,
  HASHED = 1024,
}
StateMsg = {
  [State.LOCKED]    = 'Account is locked',
  [State.VALID]     = 'Valid credentials',
  [State.DISABLED]  = 'Account is disabled',
  [State.CHANGEPW]  = 'Valid credentials, password must be changed at next logon',
  [State.PARAM]  = 'Credentials passed to script during Nmap execution',
  [State.EXPIRED]   = 'Valid credentials, account expired',
  [State.TIME_RESTRICTED] = 'Valid credentials, account cannot log in at current time',
  [State.HOST_RESTRICTED] = 'Valid credentials, account cannot log in from current host',
  [State.LOCKED_VALID]    = 'Valid credentials, account locked',
  [State.DISABLED_VALID]  = 'Valid credentials, account disabled',
  [State.HASHED]  = 'Hashed valid or invalid credentials',
}
ALL_DATA = {}
RegStorage = {
  new = function(self)
    local o = {}
    setmetatable(o, self)
    self.__index = self
    o.filter = {}
    return o
  end,
  add = function( self, tags, host, port, service, user, pass, state )
    local cred = {
      tags = tags,
      host = host,
      port = port,
      service = service,
      user = user,
      pass = pass,
      state = state
    }
    nmap.registry.creds = nmap.registry.creds or {}
    table.insert( nmap.registry.creds, cred )
  end,
  setFilter = function( self, host, port, state )
    self.filter.host = host
    self.filter.port = port
    self.filter.state = state
  end,
  getAll = function( self )
    local function get_next()
      local host, port = self.filter.host, self.filter.port
      if ( not(nmap.registry.creds) ) then return end
      for _, v in pairs(nmap.registry.creds) do
        local h = ( v.host.ip or v.host )
        if ( not(host) and not(port) ) then
          if ( not(self.filter.state) or ( v.state == self.filter.state ) ) then
            coroutine.yield(v)
          end
        elseif ( not(host) and ( port == v.port ) ) then
          if ( not(self.filter.state) or ( v.state == self.filter.state ) ) then
            coroutine.yield(v)
          end
        elseif ( ( host and ( h == host or h == host.ip ) ) and not(port) ) then
          if ( not(self.filter.state) or ( v.state == self.filter.state ) ) then
            coroutine.yield(v)
          end
        elseif ( ( host and ( h == host or h == host.ip ) ) and port.number == v.port ) then
          if ( not(self.filter.state) or ( v.state == (self.filter.state & v.state) ) ) then
            coroutine.yield(v)
          end
        end
      end
    end
    return coroutine.wrap(get_next)
  end,
}
Account = {
  new = function(self, username, password, state, port)
    local o = { username = username, password = password, state = StateMsg[state] or state, port = port}
    setmetatable(o, self)
    self.__index = self
    return o
  end,
  __tostring = function( self )
    return (
      (self.username and self.username .. ":" or "") ..
      (self.password ~= "" and self.password or "<empty>") ..
      (self.state and " - " .. self.state or "")
      )
  end,
  __lt = function (a, b)
    if a.user and b.user and a.user >= b.user then
      return false
    elseif a.pass and b.pass and a.pass >= b.pass then
      return false
    elseif a.state and b.state and a.state >= b.state then
      return false
    end
    return true
  end,
}
local function sorted_pairs (sortby)
  return function (t)
    local order = tableaux.keys(t)
    table.sort(order, sortby)
    return coroutine.wrap(function()
        for i,k in ipairs(order) do
          coroutine.yield(k, t[k])
        end
      end)
  end
end
Credentials = {
  new = function(self, tags, host, port)
    local o = {}
    setmetatable(o, self)
    self.__index = self
    o.storage = RegStorage:new()
    o.storage:setFilter(host, port)
    o.host = host
    o.port = ( port and port.number ) and port.number
    o.service = ( port and port.service ) and port.service
    if ( type(tags) ~= "table" ) then tags = {tags} end
    o.tags = tags
    return o
  end,
  add = function( self, user, pass, state )
    local pass = ( pass and #pass > 0 ) and pass or "<empty>"
    assert( self.host, "No host supplied" )
    assert( self.port, "No port supplied" )
    assert( state, "No state supplied")
    assert( self.tags, "No tags supplied")
    if ( user or pass ) then
      self.storage:add( self.tags, self.host, self.port, self.service, user, pass, state )
    end
  end,
  getCredentials = function(self, state)
    local function next_credential()
      if ( state ) then
        self.storage:setFilter(self.host, { number=self.port, service = self.service }, state)
      end
      for cred in self.storage:getAll() do
        if ( self.tags == ALL_DATA ) then
          coroutine.yield(cred)
        end
        for _,stag in pairs(self.tags) do
          for _,ctag in pairs(cred.tags) do
            if(stag == ctag) then
              coroutine.yield(cred)
            end
          end
        end
      end
      if ( state and State.PARAM == (state & State.PARAM) ) then
        local creds_global = stdnse.get_script_args('creds.global')
        local creds_service
        local creds_params
        if ( self.service ) then
          creds_service = stdnse.get_script_args('creds.' .. self.service )
        end
        if ( creds_service ) then creds_params = creds_service end
        if ( creds_global and creds_service ) then
          creds_params = creds_params .. ',' .. creds_global
        elseif ( creds_global ) then
          creds_params = creds_global
        end
        if ( not(creds_params) ) then return end
        for _, cred in ipairs(stringaux.strsplit(",", creds_params)) do
          local user, pass
          if ( cred:match(":") ) then
            user, pass = cred:match("^(.-):(.-)$")
          else
            user = cred:match("^(.*)$")
          end
          coroutine.yield( { host = self.host,
          port = self.port,
          user = user,
          pass = pass,
          state = State.PARAM,
          service = self.service } )
        end
      end
    end
    return coroutine.wrap( next_credential )
  end,
  getTable = function(self)
    local result = {}
    for v in self:getCredentials() do
      local h = ( v.host.ip or v.host )
      assert(type(h)=="string", "Could not determine a valid host")
      local svc = ("%s/%s"):format(v.port,v.service)
      result[h] = result[h] or {}
      result[h][svc] = result[h][svc] or {}
      table.insert( result[h][svc], Account:new(
          v.user ~= "" and v.user or nil,
          v.pass,
          v.state,
          v.port
          )
        )
    end
    for _, host_tbl in pairs(result) do
      for _, svc_tbl in pairs(host_tbl) do
        table.sort( svc_tbl )
      end
      setmetatable(host_tbl, {
          __pairs = sorted_pairs( function(a,b)
              return tonumber(a:match("^(%d+)")) < tonumber(b:match("^(%d+)"))
            end )
        })
    end
    setmetatable(result, {
        __pairs = sorted_pairs( function(a, b)
          return ipOps.compare_ip(a, "le", b)
        end )
      })
    local _
    if ( self.host and next(result) ) then
      _, result = next(result)
    end
    if ( self.host and self.port and next(result) ) then
      _, result = next(result)
    end
    return next(result) and result
  end,
  saveToFile = function(self, filename, fileformat)
    if ( fileformat == 'csv' ) then
      filename = filename .. '.csv'
    else
      filename = filename .. '.txt'
    end
    local f = io.open( filename, "w")
    local output = nil
    if ( not(f) ) then
      return false, ("ERROR: Failed to open file (%s)"):format(filename)
    end
    for account in self:getCredentials() do
      if ( fileformat == 'csv' ) then
        output = "\"" .. account.user .. "\",\"" .. account.pass .. "\",\"" .. StateMsg[account.state] .. "\""
      elseif ( fileformat == 'verbose') then
        output = account.user .. ":" .. account.pass .. ":" .. StateMsg[account.state]
      else
        output = account.user .. ":" .. account.pass
      end
      if ( not(f:write( output .."\n" ) ) ) then
        return false, ("ERROR: Failed to write file (%s)"):format(filename)
      end
    end
    f:close()
    return true
  end,
  __tostring = function(self)
    local all = self:getTable()
    if ( all ) then return tostring(all) end
  end,
}
return _ENV;
