local coroutine = require "coroutine"
local creds = require "creds"
local io = require "io"
local nmap = require "nmap"
local os = require "os"
local stdnse = require "stdnse"
local table = require "table"
local unpwdb = require "unpwdb"
local math = require "math"
_ENV = stdnse.module("brute", stdnse.seeall)
Options = {
  new = function (self)
    local o = {}
    setmetatable(o, self)
    self.__index = self
    o.emptypass = self.checkBoolArg("brute.emptypass", false)
    o.useraspass = self.checkBoolArg("brute.useraspass", false)
    o.firstonly = self.checkBoolArg("brute.firstonly", false)
    o.passonly = self.checkBoolArg("brute.passonly", false)
    o.killstagnated = self.checkBoolArg("brute.killstagnated", true)
    o.max_retries = tonumber(stdnse.get_script_args("brute.retries")) or 2
    o.delay = tonumber(stdnse.get_script_args("brute.delay")) or 0
    o.max_guesses = tonumber(stdnse.get_script_args("brute.guesses")) or 0
    return o
  end,
  checkBoolArg = function (arg, default)
    local val = stdnse.get_script_args(arg) or default
    return (val == "true" or val == true or tonumber(val) == 1)
  end,
  setMode = function (self, mode)
    local modes = {
      "password",
      "user",
      "creds",
    }
    local supported = false
    for _, m in ipairs(modes) do
      if mode == m then
        supported = true
      end
    end
    if not supported then
      stdnse.debug1("ERROR: brute.options.setMode: mode %s not supported", mode)
      return false, "Unsupported mode"
    else
      self.mode = mode
    end
    return true
  end,
  setOption = function (self, param, value)
    self[param] = value
  end,
  setTitle = function (self, title)
    self.title = title
  end,
}
Error = {
  retry = false,
  new = function (self, msg)
    local o = {
      msg = msg,
      done = false,
      reduce = nil,
    }
    setmetatable(o, self)
    self.__index = self
    return o
  end,
  isRetry = function (self)
    return self.retry
  end,
  setRetry = function (self, r)
    self.retry = r
  end,
  setAbort = function (self, b)
    self.abort = b
  end,
  isAbort = function (self)
    return self.abort
  end,
  getMessage = function (self)
    return self.msg
  end,
  isDone = function (self)
    return self.done
  end,
  setDone = function (self, b)
    self.done = b
  end,
  setInvalidAccount = function (self, username)
    self.invalid_account = username
  end,
  isInvalidAccount = function (self)
    return self.invalid_account
  end,
  setReduce = function (self, r)
    self.reduce = r
  end,
  isReduce = function (self)
    if self.reduce then
      return true
    end
    return false
  end,
}
Batch = {
  new = function (self, lim, stime)
    local o = {
      limit = lim or 3, -- maximum number of items
      full = false,
      data = {}, -- storage
      size = 0, -- current number of items
      start_time = stime or 0,
    }
    setmetatable(o, self)
    self.__index = self
    return o
  end,
  add = function (self, obj)
    if self.size < self.limit then
      self.data[self.size + 1] = obj
      self.size = self.size + 1
      return true
    end
    return false
  end,
  isFull = function (self)
    if self.size >= self.limit then
      return true
    end
    return false
  end,
  getData = function (self)
    return self.data
  end,
  getSize = function (self)
    return self.size
  end,
  getStartTime = function (self)
    return self.start_time
  end,
  getLimit = function (self)
    return self.limit
  end,
  setLimit = function (self, lim)
    self.limit = lim
  end,
}
Engine = {
  STAT_INTERVAL = 20,
  THREAD_TO_ENGINE = {},
  new = function (self, driver, host, port, options)
    if getmetatable(Engine.THREAD_TO_ENGINE) == nil then
      setmetatable(Engine.THREAD_TO_ENGINE, {
          __mode = "k",
        })
    end
    local o = {
      driver = driver,
      host = host,
      port = port,
      driver_options = options,
      terminate_all = false,
      error = nil,
      counter = 0,
      threads = {},
      tps = {},
      iterator = nil,
      usernames = usernames_iterator(),
      passwords = passwords_iterator(),
      found_accounts = {},
      account_guesses = {},
      options = Options:new(),
      retry_accounts = {},
      initial_accounts_exhausted = false,
      batch = nil,
      tick = 0,
    }
    setmetatable(o, self)
    self.__index = self
    o.max_threads = tonumber(stdnse.get_script_args "brute.threads") or 20
    o.start_threads = tonumber(stdnse.get_script_args "brute.start") or 5
    return o
  end,
  setUsernameIterator = function (self, usernameIterator)
    self.usernames = usernameIterator
  end,
  setPasswordIterator = function (self, passwordIterator)
    self.passwords = passwordIterator
  end,
  setMaxThreads = function (self, max)
    self.max_threads = max
  end,
  threadCount = function (self)
    local count = 0
    for thread in pairs(self.threads) do
      if coroutine.status(thread) == "dead" then
        self.threads[thread] = nil
      else
        count = count + 1
      end
    end
    return count
  end,
  activeThreads = function (self)
    local count = 0
    for thread, v in pairs(self.threads) do
      if v.guesses ~= nil then
        count = count + 1
      end
    end
    return count
  end,
  get_next_credential = function (self)
    local function next_credential ()
      for user, pass in self.iterator do
        self.used_creds = self.used_creds or {}
        pass = pass or "nil"
        if not self.used_creds[user .. pass] then
          self.used_creds[user .. pass] = true
          coroutine.yield(user, pass)
        end
      end
      while true do
        coroutine.yield(nil, nil)
      end
    end
    return coroutine.wrap(next_credential)
  end,
  doAuthenticate = function (self)
    local status, response
    local next_credential = self:get_next_credential()
    local tries = self.options.max_retries + 1
    local username, password
    local thread_data = Engine.getThreadData(coroutine.running())
    assert(thread_data, "Unknown coroutine is running")
    repeat
      local driver = self.driver:new(self.host, self.port, self.driver_options)
      status, response = driver:connect()
      if not status then
        if response and not response.isReduce then
          response = Error:new("Connect error: " .. response)
          response:setRetry(true)
        end
        if response:isReduce() then
          local ret_creds = {}
          ret_creds.connect_phase = true
          return false, response, ret_creds
        end
      else
        if not username and not password then
          repeat
            if #self.retry_accounts > 0 then
              username = self.retry_accounts[#self.retry_accounts].username
              password = self.retry_accounts[#self.retry_accounts].password
              table.remove(self.retry_accounts, #self.retry_accounts)
            else
              username, password = next_credential()
            end
            thread_data.username = username
            thread_data.password = password
            if not username and not password then
              driver:disconnect()
              self.initial_accounts_exhausted = true
              return false
            end
          until (not self.found_accounts or not self.found_accounts[username])
            and (self.options.max_guesses == 0 or not self.account_guesses[username]
                  or self.options.max_guesses > self.account_guesses[username])
          self.account_guesses[username] = self.account_guesses[username]
                                           and self.account_guesses[username] + 1 or 1
        end
        if Engine.terminate_all then
          driver:disconnect()
          driver = nil
          return false
        end
        local c
        if username and #username > 0 then
          c = ("%s/%s"):format(username, #password > 0 and password or "<empty>")
        else
          c = ("%s"):format(#password > 0 and password or "<empty>")
        end
        local msg = (tries <= self.options.max_retries) and "Re-trying" or "Trying"
        stdnse.debug2("%s %s against %s:%d", msg, c, self.host.ip, self.port.number)
        status, response = driver:login(username, password)
        driver:disconnect()
        driver = nil
        if not status and response:isReduce() then
          local ret_creds = {}
          ret_creds.username = username
          ret_creds.password = password
          return false, response, ret_creds
        end
      end
      tries = tries - 1
    until status or (response and not (response:isRetry())) or tries <= 0
    self.counter = self.counter + 1
    return status, response
  end,
  login = function (self, cvar)
    local condvar = nmap.condvar(cvar)
    local thread_data = self.threads[coroutine.running()]
    local interval_start = os.time()
    while true do
      if (self.terminate_all or thread_data.terminate)
        or (self.initial_accounts_exhausted and #self.retry_accounts == 0) then
        break
      end
      self.tick = self.tick + 1
      if not (self.batch:isFull()) and not thread_data.in_batch then
        self.batch:add(coroutine.running())
        thread_data.in_batch = true
        thread_data.ready = false
      end
      local status, response, ret_creds = self:doAuthenticate()
      if thread_data.in_batch then
        thread_data.ready = true
      end
      if status then
        if not self.found_accounts or self.found_accounts[response.username] == nil then
          if not self.options.nostore then
            local c = creds.Credentials:new(self.options.script_name, self.host, self.port)
            c:add(response.username, response.password, response.state)
          else
            self.credstore = self.credstore or {}
            table.insert(self.credstore, tostring(response))
          end
          stdnse.debug1("Discovered account: %s", tostring(response))
          if not self.options.passonly then
            self.found_accounts[response.username] = true
          end
          if self.options.firstonly then
            self.terminate_all = true
          end
        end
      elseif ret_creds then
        if not ret_creds.connect_phase then
            self.retry_accounts[#self.retry_accounts + 1] = {
              username = ret_creds.username,
              password = ret_creds.password,
            }
        end
        thread_data.protocol_error = true
        condvar "signal"
        condvar "wait"
      else
        if response and response:isAbort() then
          self.terminate_all = true
          self.error = response:getMessage()
          break
        elseif response and response:isDone() then
          break
        elseif response and response:isInvalidAccount() then
          self.found_accounts[response:isInvalidAccount()] = true
        end
      end
      local timediff = (os.time() - interval_start)
      thread_data.guesses = (thread_data.guesses and thread_data.guesses + 1 or 1)
      if timediff > Engine.STAT_INTERVAL then
        interval_start = os.time()
        local tps = self.counter / (os.time() - self.starttime)
        table.insert(self.tps, tps)
        stdnse.debug2("threads=%d,tps=%.1f", self:activeThreads(), tps)
      end
      if self.options.delay > 0 then
        stdnse.sleep(self.options.delay)
      end
      condvar "signal"
    end
    condvar "signal"
  end,
  addWorker = function (self, cvar)
    local co = stdnse.new_thread(self.login, self, cvar)
    Engine.THREAD_TO_ENGINE[co] = self
    self.threads[co] = {
      running = true,
      protocol_error = nil,
      attempt = 0,
      in_batch = false,
      ready = false,
      connection_error = nil,
      con_error_reason = nil,
      username = nil,
      password = nil,
    }
    return co
  end,
  addWorkerN = function (self, cvar, n)
    assert(n >= 0)
    for i = 1, n do
      self:addWorker(cvar)
    end
  end,
  renewBatch = function (self)
    if self.batch then
      local size = self.batch:getSize()
      local data = self.batch:getData()
      for i = 1, size do
        if self.threads[data[i]] then
          self.threads[data[i]].in_batch = false
          self.threads[data[i]].ready = false
        end
      end
    end
    self.batch = Batch:new(math.min(self:threadCount(), 3), self.tick)
  end,
  readyBatch = function (self)
    if not self.batch then
      return false
    end
    local n = self.batch:getSize()
    local data = self.batch:getData()
    if n == 0 then
      return false
    end
    for i = 1, n do
      if self.threads[data[i]] and coroutine.status(data[i]) ~= "dead" and self.threads[data[i]].in_batch then
        if not self.threads[data[i]].ready then
          return false
        end
      end
    end
    return true
  end,
  start = function (self)
    local cvar = {}
    local condvar = nmap.condvar(cvar)
    assert(self.options.script_name, "SCRIPT_NAME was not set in options.script_name")
    assert(self.port.number and self.port.protocol, "Invalid port table detected")
    self.port.service = self.port.service or "unknown"
    if self.driver:new(self.host, self.port, self.driver_options).check then
      local status, response = self.driver:new(self.host, self.port, self.driver_options):check()
      if not status then
        return false, response
      end
    end
    local usernames = self.usernames
    local passwords = self.passwords
    if "function" ~= type(usernames) then
      return false, "Invalid usernames iterator"
    end
    if "function" ~= type(passwords) then
      return false, "Invalid passwords iterator"
    end
    local mode = self.options.mode or stdnse.get_script_args "brute.mode"
    if not mode and stdnse.get_script_args "brute.credfile" then
      if stdnse.get_script_args "userdb" or stdnse.get_script_args "passdb" then
        return false, "\n  ERROR: brute.credfile can't be used in combination with userdb/passdb"
      end
      mode = 'creds'
    end
    if self.options.passonly then
      local function single_user_iter (next)
        local function next_user ()
          coroutine.yield ""
        end
        return coroutine.wrap(next_user)
      end
      if self.iterator == nil then
        self.iterator = Iterators.user_pw_iterator(single_user_iter(), passwords)
      end
    elseif mode == 'creds' then
      local credfile = stdnse.get_script_args "brute.credfile"
      if not credfile then
        return false, "No credential file specified (see brute.credfile)"
      end
      local f = io.open(credfile, "r")
      if not f then
        return false, ("Failed to open credfile (%s)"):format(credfile)
      end
      self.iterator = Iterators.credential_iterator(f)
    elseif mode and mode == 'user' then
      self.iterator = self.iterator or Iterators.user_pw_iterator(usernames, passwords)
    elseif mode and mode == 'pass' then
      self.iterator = self.iterator or Iterators.pw_user_iterator(usernames, passwords)
    elseif mode then
      return false, ("Unsupported mode: %s"):format(mode)
    elseif self.iterator == nil then
      self.iterator = Iterators.pw_user_iterator(usernames, passwords)
    end
    if (not mode or mode == 'user' or mode == 'pass') and self.options.useraspass then
      if not self.options.passonly then
        self.iterator = unpwdb.concat_iterators(
          Iterators.pw_same_as_user_iterator(usernames, "lower"),
          self.iterator
        )
      end
    end
    if (not mode or mode == 'user' or mode == 'pass') and self.options.emptypass then
      local function empty_pass_iter ()
        local function next_pass ()
          coroutine.yield ""
        end
        return coroutine.wrap(next_pass)
      end
      self.iterator = Iterators.account_iterator(usernames, empty_pass_iter(), mode or "pass")
    end
    self.starttime = os.time()
    local start_threads = self.start_threads
    if nmap.socket.get_stats().connect_waiting > 0 then
      start_threads = 1
    end
    self:addWorkerN(cvar, start_threads)
    self:renewBatch()
    local revive = false
    local killed_one = false
    local error_since_batch_start = false
    local stagnation_count = 0 -- number of times when all threads are stopped because of exceptions
    local quick_start = true
    local stagnated = true
    while true do
      local thread_count = self:threadCount()
      if thread_count <= 0 then
        if (self.initial_accounts_exhausted and #self.retry_accounts == 0) or self.terminate_all then
          break
        else
          revive = true
        end
      end
      killed_one = false
      error_since_batch_start = false
      stagnated = true
      for co, v in pairs(self.threads) do
        if not v.connection_error then
          stagnated = false
        end
        if v.protocol_error or v.connection_error then
          if v.attempt >= self.batch:getStartTime() then
            error_since_batch_start = true
          end
          if not killed_one then
            v.terminate = true
            killed_one = true
            if v.protocol_error then
              stdnse.debug2("Killed one thread because of PROTOCOL exception")
            else
              stdnse.debug2("Killed one thread because of CONNECTION exception")
            end
          end
          v.protocol_error = nil
          v.connection_error = nil
        else
        end
      end
      if stagnated == true then
        stagnation_count = stagnation_count + 1
        if stagnation_count == 100 then
          stdnse.debug1("WARNING: The service seems to have failed or is heavily firewalled... Consider aborting.")
          if self.options.killstagnated then
            self.error = "The service seems to have failed or is heavily firewalled..."
            self.terminate_all = true
          end
          stagnation_count = 0
        end
      else
        stagnation_count = 0
      end
      if killed_one then
        quick_start = false
      end
      if not killed_one then
        local waiting = nmap.socket.get_stats().connect_waiting
        if waiting ~= 0 then
          local kill_count = 1
          if waiting > 5 then
            kill_count = math.max(math.floor(thread_count / 2), 1)
          end
          for co, v in pairs(self.threads) do
            if coroutine.status(co) ~= "dead" then
              stdnse.debug2("Killed one because of RESOURCE management")
              v.terminate = true
              killed_one = true
              kill_count = kill_count - 1
              if kill_count == 0 then
                break
              end
            end
          end
        end
      end
      if error_since_batch_start
        or not killed_one and thread_count < self.batch:getLimit()
        or (thread_count > 0 and self.tick - self.batch:getStartTime() > 10) then
        self:renewBatch()
      end
      if (not killed_one and self.batch:isFull() and thread_count < self.max_threads)
        or revive then
        local num_to_add = 1
        if quick_start then
          num_to_add = math.min(self.max_threads - thread_count, thread_count)
        end
        self:addWorkerN(cvar, num_to_add)
        self:renewBatch()
        revive = false
      end
      local threads = self:threadCount()
      stdnse.debug2("Status: #threads = %d, #retry_accounts = %d, initial_accounts_exhausted = %s, waiting = %d",
        threads, #self.retry_accounts, tostring(self.initial_accounts_exhausted),
        nmap.socket.get_stats().connect_waiting)
      if threads > 0 then
        condvar "broadcast"
        condvar "wait"
      end
    end
    local valid_accounts
    if not self.options.nostore then
      valid_accounts = creds.Credentials:new(self.options.script_name, self.host, self.port):getTable()
    else
      valid_accounts = self.credstore
    end
    local result = stdnse.output_table()
    if valid_accounts and #valid_accounts > 0 then
      result[self.options.title or "Accounts"] = valid_accounts
    else
      result.Accounts = "No valid accounts found"
    end
    local sum = 0
    for _, v in ipairs(self.tps) do
      sum = sum + v
    end
    local time_diff = (os.time() - self.starttime)
    time_diff = (time_diff == 0) and 1 or time_diff
    local tps = (sum == 0) and (self.counter / time_diff) or (sum / #self.tps)
    result.Statistics = ("Performed %d guesses in %d seconds, average tps: %.1f"):format( self.counter, time_diff, tps )
    if self.options.max_guesses > 0 then
      for user, guesses in pairs(self.account_guesses) do
        if guesses == self.options.max_guesses then
          result.Information = ("Guesses restricted to %d tries per account to avoid lockout"):format(self.options.max_guesses)
          break
        end
      end
    end
    if self.error then
      result.ERROR = self.error
      return false, result
    end
    return true, result
  end,
  getEngine = function (co)
    local engine = Engine.THREAD_TO_ENGINE[co]
    if not engine then
      stdnse.debug1("WARNING: No engine associated with %s", coroutine.running())
    end
    return engine
  end,
  getThreadData = function (co)
    local engine = Engine.getEngine(co)
    if not engine then
      return nil
    end
    return engine.threads[co]
  end,
}
function usernames_iterator ()
  local status, usernames = unpwdb.usernames()
  if not status then
    return "Failed to load usernames"
  end
  return usernames
end
function passwords_iterator ()
  local status, passwords = unpwdb.passwords()
  if not status then
    return "Failed to load passwords"
  end
  return passwords
end
Iterators = {
  account_iterator = function (users, pass, mode)
    local function next_credential ()
      local outer, inner
      if "table" == type(users) then
        users = unpwdb.table_iterator(users)
      end
      if "table" == type(pass) then
        pass = unpwdb.table_iterator(pass)
      end
      if mode == 'pass' then
        outer, inner = pass, users
      elseif mode == 'user' then
        outer, inner = users, pass
      else
        return
      end
      for o in outer do
        for i in inner do
          if mode == 'pass' then
            coroutine.yield(i, o)
          else
            coroutine.yield(o, i)
          end
        end
        inner "reset"
      end
      while true do
        coroutine.yield(nil, nil)
      end
    end
    return coroutine.wrap(next_credential)
  end,
  user_pw_iterator = function (users, pass)
    return Iterators.account_iterator(users, pass, "user")
  end,
  pw_user_iterator = function (users, pass)
    return Iterators.account_iterator(users, pass, "pass")
  end,
  pw_same_as_user_iterator = function (users, case)
    local function next_credential ()
      for user in users do
        if case == 'upper' then
          coroutine.yield(user, user:upper())
        elseif case == 'lower' then
          coroutine.yield(user, user:lower())
        else
          coroutine.yield(user, user)
        end
      end
      users "reset"
      while true do
        coroutine.yield(nil, nil)
      end
    end
    return coroutine.wrap(next_credential)
  end,
  pw_ucase_iterator = function (users, passwords, mode)
    local function next_credential ()
      for user, pass in Iterators.account_iterator(users, passwords, mode) do
        coroutine.yield(user, pass:upper())
      end
      while true do
        coroutine.yield(nil, nil)
      end
    end
    return coroutine.wrap(next_credential)
  end,
  credential_iterator = function (f)
    local function next_credential ()
      local c = {}
      for line in f:lines() do
        if not (line:match "^#!comment:") then
          local trim = function (s)
            return s:match '^()%s*$' and '' or s:match '^%s*(.*%S)'
          end
          line = trim(line)
          local user, pass = line:match "^([^%/]*)%/(.*)$"
          coroutine.yield(user, pass)
        end
      end
      f:close()
      while true do
        coroutine.yield(nil, nil)
      end
    end
    return coroutine.wrap(next_credential)
  end,
  unpwdb_iterator = function (mode)
    local status, users, passwords
    status, users = unpwdb.usernames()
    if not status then
      return
    end
    status, passwords = unpwdb.passwords()
    if not status then
      return
    end
    return Iterators.account_iterator(users, passwords, mode)
  end,
}
checkwrap = {
  connect = true,
  send = true,
  receive = true,
  receive_lines = true,
  receive_buf = true,
  receive_bytes = true,
}
BruteSocket = {
  new = function (self)
    local o = {
      socket = nil,
    }
    setmetatable(o, self)
    self.__index = function (instance, key)
      local f = rawget(self, key)
      if f then
        return f
      else
        f = instance.socket[key]
      end
      if checkwrap[key] then
        return function(s, ...)
          local status, err = f(instance.socket, ...)
          instance:checkStatus(status, err)
          return status, err
        end
      elseif type(f) == "function" then
        return function (s, ...)
          return f(instance.socket, ...)
        end
      end
      return f
    end
    o.socket = nmap.new_socket()
    return o
  end,
  getSocket = function (self)
    return self.socket
  end,
  checkStatus = function (self, status, err)
    if not status and (err == "ERROR" or err == "TIMEOUT") then
      local engine = Engine.getEngine(coroutine.running())
      if not engine then
        stdnse.debug2("WARNING: No associated engine detected for %s", coroutine.running())
        return -- behave like a usual socket
      end
      local thread_data = Engine.getThreadData(coroutine.running())
      engine.retry_accounts[#engine.retry_accounts + 1] = {
        username = thread_data.username,
        password = thread_data.password,
      }
      thread_data.connection_error = true
      thread_data.con_error_reason = err
    end
  end,
}
function new_socket ()
  return BruteSocket:new()
end
return _ENV
