local io = require "io"
local nmap = require "nmap"
local os = require "os"
local stdnse = require "stdnse"
local datetime = require "datetime"
_ENV = stdnse.module("unpwdb", stdnse.seeall)
local usertable = {}
local passtable = {}
local customdata = false
local args = nmap.registry.args
local userfile = function()
  if args.userdb then
    customdata = true
    return args.userdb
  end
  return nmap.fetchfile("nselib/data/usernames.lst")
end
local passfile = function()
  if args.passdb then
    customdata = true
    return args.passdb
  end
  return nmap.fetchfile("nselib/data/passwords.lst")
end
local filltable = function(filename, table)
  if #table ~= 0 then
    return true
  end
  local file = io.open(filename, "r")
  if not file then
    return false
  end
  for l in file:lines() do
    if not l:match("#!comment:") then
      table[#table + 1] = l
    end
  end
  file:close()
  return true
end
table_iterator = function(table)
  local i = 1
  return function(cmd)
    if cmd == "reset" then
      i = 1
      return
    end
    local elem = table[i]
    if elem then i = i + 1 end
    return elem
  end
end
timelimit = function()
  if args.notimelimit then
    return nil
  end
  if args["unpwdb.timelimit"] then
    local limit, err = stdnse.parse_timespec(args["unpwdb.timelimit"])
    if not limit then
      error(err)
    end
    return limit
  end
  local t = nmap.timing_level()
  if t <= 3 then
    return (customdata and 900) or 600
  elseif t == 4 then
    return (customdata and 450) or 300
  elseif t == 5 then
    return (customdata and 270) or 180
  end
end
local usernames_raw = function()
  local path = userfile()
  if not path then
    return false, "Cannot find username list"
  end
  if not filltable(path, usertable) then
    return false, "Error parsing username list"
  end
  return true, table_iterator(usertable)
end
local passwords_raw = function()
  local path = passfile()
  if not path then
    return false, "Cannot find password list"
  end
  if not filltable(path, passtable) then
    return false, "Error parsing password list"
  end
  return true, table_iterator(passtable)
end
limited_iterator = function(iterator, time_limit, count_limit, label)
  time_limit = (time_limit and time_limit > 0) and time_limit
  count_limit = (count_limit and count_limit > 0) and count_limit
  local start = os.time()
  local count = 0
  label = label or "limited_iterator"
  return function(cmd)
    if cmd == "reset" then
      count = 0
    else
      count = count + 1
    end
    if count_limit and count > count_limit then
      stdnse.verbose1("%s: Count limit %d exceeded.", label, count_limit)
      return
    end
    if time_limit and os.time() - start >= time_limit then
      stdnse.verbose1("%s: Time limit %s exceeded.", label, datetime.format_time(time_limit))
      return
    end
    return iterator(cmd)
  end
end
usernames = function(time_limit, count_limit)
  local status, iterator
  status, iterator = usernames_raw()
  if not status then
    return false, iterator
  end
  time_limit = time_limit or timelimit()
  if not count_limit and args["unpwdb.userlimit"] then
    count_limit = tonumber(args["unpwdb.userlimit"])
  end
  return true, limited_iterator(iterator, time_limit, count_limit, "usernames")
end
passwords = function(time_limit, count_limit)
  local status, iterator
  status, iterator = passwords_raw()
  if not status then
    return false, iterator
  end
  time_limit = time_limit or timelimit()
  if not count_limit and args["unpwdb.passlimit"] then
    count_limit = tonumber(args["unpwdb.passlimit"])
  end
  return true, limited_iterator(iterator, time_limit, count_limit, "passwords")
end
function concat_iterators (iter1, iter2)
  local function helper (next_iterator, command, first, ...)
    if first ~= nil then
      return first, ...
    elseif next_iterator ~= nil then
      return helper(nil, command, next_iterator(command))
    end
  end
  local function iterator (command)
    if command == "reset" then
      iter1 "reset"
      iter2 "reset"
    else
      return helper(iter2, command, iter1(command))
    end
  end
  return iterator
end
function filter_iterator (iterator, filter)
  return function (command)
    if command == "reset" then
      iterator "reset"
    else
      local val = iterator(command)
      while val and not filter(val) do
        val = iterator(command)
      end
      return val
    end
  end
end
return _ENV;
