--- The credential class stores found credentials in the Nmap registry
--
-- The credentials library may be used by scripts to store credentials in
-- a common format in the nmap registry. The Credentials class serves as
-- a primary interface for scripts to the library.
--
-- The State table keeps track of possible account states and a corresponding
-- message to return for each state.
--
-- The following code illustrates how a script may add discovered credentials
-- to the database:
-- <code>
-- local c = creds.Credentials:new( SCRIPT_NAME, host, port )
-- c:add("patrik", "secret", creds.State.VALID )
-- </code>
--
-- The following code illustrates how a script can return a table of discovered
-- credentials at the end of execution:
-- <code>
-- return tostring(creds.Credentials:new(SCRIPT_NAME, host, port))
-- </code>
--
-- The following code illustrates how a script may iterate over discovered
-- credentials:
-- <code>
-- local c = creds.Credentials:new(creds.ALL_DATA, host, port)
-- for cred in c:getCredentials(creds.State.VALID) do
--   showContentForUser(cred.user, cred.pass)
-- end
-- </code>
--
-- The library also enables users to add credentials through script arguments
-- either globally or per service. These credentials may be retrieved by script
-- through the same functions as any other discovered credentials. Arguments
-- passed using script arguments will be added with the PARAM state. The
-- following code may be used by a scripts to retrieve these credentials:
-- <code>
-- local c = creds.Credentials:new(creds.ALL_DATA, host, port)
-- for cred in c:getCredentials(creds.State.PARAM) do
--   ... do something ...
-- end
-- </code>
--
-- Any globally added credentials will be made available to all scripts,
-- regardless of what service is being filtered through the host and port
-- arguments when instantiating the Credentials class. Service specific
-- arguments will only be made available to scripts with ports matching
-- the service name. The following two examples illustrate how credentials are
-- added globally and for the http service:
-- <code>
-- --script-args creds.global='admin:nimda'
-- --script-args creds.http='webadmin:password'
-- </code>
--
-- The service name at this point may be anything and the entry is created
-- dynamically without validating whether the service exists or not.
--
-- The credential argument is not documented in this library using the <at>args
-- function as the argument would incorrectly show up in all scripts making use
-- of this library. This would show that credentials could be added to scripts
-- that do not make use of this function. Therefore any scripts that make use
-- of the credentials passing arguments need to have appropriate documentation
-- added to them.
--
--
-- The following code illustrates how a script may save its discovered credentials
-- to a file:
-- <code>
-- local c = creds.Credentials:new( SCRIPT_NAME, host, port )
-- c:add("patrik", "secret", creds.State.VALID )
-- status, err = c:saveToFile("outputname","csv")
-- </code>
--
-- Supported output formats are CSV, verbose and plain.  In both verbose and plain
-- records are separated by colons.  The difference between the two is that verbose
-- includes the credential state.  The file extension is automatically added to
-- the filename based on the type requested.
--
-- @author "Patrik Karlsson <patrik@cqure.net>"
-- @copyright Same as Nmap--See http://nmap.org/book/man-legal.html

-- Version 0.4
-- Created 2011/02/06 - v0.1 - created by Patrik Karlsson <patrik@cqure.net>
-- Revised 2011/27/06 - v0.2 - revised by Patrik Karlsson <patrik@cqure.net>
--                * added documentation
--                * added getCredentials function
--
-- Revised 2011/05/07 - v0.3 - revised by Patrik Karlsson <patrik@cqure.net>
--                * modified getCredentials to return an iterator
--                * added support for adding credentials as
--                  script arguments
--
-- Revised 2011/09/04 - v0.4 - revised by Tom Sellers
--                * added saveToFile function for saving credential
--                * table to file in CSV or text formats

local bit = require "bit"
local coroutine = require "coroutine"
local io = require "io"
local ipOps = require "ipOps"
local nmap = require "nmap"
local stdnse = require "stdnse"
local table = require "table"
_ENV = stdnse.module("creds", stdnse.seeall)


--- Table mapping the different account states to their number
--
-- Also available is the <code>StateMsg</code> table, used to map these numbers
-- to a description.
-- @class table
-- @name State
-- @field LOCKED Account is locked
-- @field VALID Valid credentials
-- @field DISABLED Account is disabled
-- @field CHANGEPW Valid credentials, password must be changed at next logon
-- @field PARAM Credentials passed to script during Nmap execution
-- @field EXPIRED Valid credentials, account expired
-- @field TIME_RESTRICTED Valid credentials, account cannot log in at current time
-- @field HOST_RESTRICTED Valid credentials, account cannot log in from current host
-- @field LOCKED_VALID Valid credentials, account locked
-- @field DISABLED_VALID Valid credentials, account disabled
-- @field HASHED Hashed valid or invalid credentials
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 = "all_script_data"

-- The RegStorage class
RegStorage = {

  --- Creates a new RegStorage instance
  --
  -- @return a new instance
  new = function(self)
    local o = {}
    setmetatable(o, self)
    self.__index = self
    o.filter = {}
    return o
  end,

  --- Add credentials to storage
  --
  -- @param scriptname the name of the script adding the credentials
  -- @param host host table, name or ip
  -- @param port number containing the port of the service
  -- @param service the name of the service
  -- @param user the name of the user
  -- @param pass the password of the user
  -- @param state of the account
  add = function( self, scriptname, host, port, service, user, pass, state )
    local cred = {
      scriptname = scriptname,
      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,

  --- Sets the storage filter
  --
  -- @param host table containing the host
  -- @param port table containing the port
  -- @param state table containing the account state
  setFilter = function( self, host, port, state )
    self.filter.host = host
    self.filter.port = port
    self.filter.state = state
  end,

  --- Returns a credential iterator matching the selected filters
  --
  -- @return a credential iterator
  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 == bit.band(self.filter.state, v.state) ) ) then
            coroutine.yield(v)
          end
        end
      end
    end
    return coroutine.wrap(get_next)
  end,

}

-- The credentials class
Credentials = {

  --- Creates a new instance of the Credentials class
  -- @param scriptname string containing the name of the script
  -- @param host table as received by the scripts action method
  -- @param port table as received by the scripts action method
  new = function(self, scriptname, 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
    o.scriptname = scriptname
    return o
  end,

  --- Add a discovered credential
  --
  -- @param user the name of the user
  -- @param pass the password of the user
  -- @param state of the account
  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.scriptname, "No scriptname supplied")

    -- there are cases where we will only get a user or password
    -- so as long we have one of them, we're good
    if ( user or pass ) then
      self.storage:add( self.scriptname, self.host, self.port, self.service, user, pass, state )
    end
  end,

  --- Returns a credential iterator
  --
  -- @see State
  -- @param state mask containing values from the <Code>State</code> table
  -- @return credential iterator, returning a credential each time it's
  --         called. Unless filtered by the state mask all credentials
  --         for the host, port match are iterated over.
  --         The credential table has the following fields:
  --         <code>host</code> - table as received by the action function
  --         <code>port</code> - number containing the port number
  --         <code>user</code> - string containing the user name
  --         <code>pass</code> - string containing the user password
  --         <code>state</code> - a state number
  --         <code>service</code> - string containing the name of the service
  --         <code>scriptname</code> - string containing the name of the
  --                                   script that added the credential
  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.scriptname == ALL_DATA ) or
          ( cred.scriptname == self.scriptname ) ) then
          coroutine.yield(cred)
        end
      end

      if ( state and State.PARAM == bit.band(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(stdnse.strsplit(",", creds_params)) do
          -- if the credential contains a ':' we have a user + pass pair
          -- if not, we only have a user with an empty password
          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,

  --- Returns a table of credentials
  --
  -- @return tbl table containing the discovered credentials
  getTable = function(self)
    local result = {}

    for v in self.storage:getAll() do
      local h = ( v.host.ip or v.host )
      local svc = ("%s/%s"):format(v.port,v.service)
      local c
      if ( v.user and #v.user > 0 ) then
        if StateMsg[v.state] then
          c = ("%s:%s - %s"):format(v.user, v.pass, StateMsg[v.state])
        else
          c = ("%s:%s"):format(v.user, v.pass)
        end
      else
        if StateMsg[v.state] then
          c = ("%s - %s"):format(v.pass, StateMsg[v.state])
        else
          c = ("%s"):format(v.pass)
        end
      end
      local script = v.scriptname
      assert(type(h)=="string", "Could not determine a valid host")

      if ( script == self.scriptname or self.scriptname == ALL_DATA ) then
        result[h] = result[h] or {}
        result[h][svc] = result[h][svc] or {}
        table.insert( result[h][svc], c )
      end
    end

    local output = {}
    for hostname, host in pairs(result) do
      local host_tbl = { name = hostname }
      for svcname, service in pairs(host) do
        local svc_tbl = { name = svcname }
        for _, account in ipairs(service) do
          table.insert(svc_tbl, account)
        end
        -- sort the accounts
        table.sort( svc_tbl, function(a,b) return a<b end)
        table.insert( host_tbl, svc_tbl )
      end
      -- sort the services
      table.sort( host_tbl,
      function(a,b)
        return tonumber(a.name:match("^(%d+)")) < tonumber(b.name:match("^(%d+)"))
      end
      )
      table.insert( output, host_tbl )
    end

    -- sort the IP addresses
    table.sort( output, function(a, b) return ipOps.compare_ip(a.name, "le", b.name) end )
    if ( self.host and self.port and #output > 0 ) then
      output = output[1][1]
      output.name = nil
    elseif ( self.host and #output > 0 ) then
      output = output[1]
      output.name = nil
    end
    return (#output > 0 ) and output
  end,

  -- Saves credentials in the current object to file
  -- @param filename string name of the file
  -- @param fileformat string file format type, values = csv | verbose | plain (default)
  -- @return status true on success, false on failure
  -- @return err string containing the error if status is false
  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,

  --- Get credentials with optional host and port filter
  -- If no filters are supplied all records are returned
  --
  -- @param host table or string containing the host to filter
  -- @param port number containing the port to filter
  -- @return table suitable from <code>stdnse.format_output</code>
  __tostring = function(self)
    local all = self:getTable()
    if ( all ) then return stdnse.format_output(true, all) end
  end,

}

return _ENV;