module Msf::Exploit::Retry
  # Retry the block until it returns a truthy value. Each iteration attempt will
  # be performed with an exponential backoff. If the timeout period surpasses,
  # nil is returned.
  #
  # @param Integer timeout the number of seconds to wait before the operation times out
  # @return the truthy value of the block is returned or nil if it timed out
  def retry_until_truthy(timeout:)
    start_time = Process.clock_gettime(Process::CLOCK_MONOTONIC, :second)
    ending_time = start_time + timeout
    retry_count = 0
    while Process.clock_gettime(Process::CLOCK_MONOTONIC, :second) < ending_time
      result = yield
      return result if result

      retry_count += 1
      remaining_time_budget = ending_time - Process.clock_gettime(Process::CLOCK_MONOTONIC, :second)
      break if remaining_time_budget <= 0

      delay = 2**retry_count
      if delay >= remaining_time_budget
        delay = remaining_time_budget
        vprint_status("Final attempt. Sleeping for the remaining #{delay} seconds out of total timeout #{timeout}")
      else
        vprint_status("Sleeping for #{delay} seconds before attempting again")
      end

      sleep delay
    end

    nil
  end

  # Poll the block at a fixed interval until it returns a truthy value or the
  # timeout expires. Unlike retry_until_truthy, this uses a consistent delay
  # between attempts rather than exponential backoff — useful when checking for
  # a state change that could happen at any moment (e.g. waiting for a session).
  #
  # @param timeout [Integer] maximum seconds to wait
  # @param interval [Numeric] seconds between each poll (default: 1)
  # @return the truthy value of the block, or nil if timed out
  def poll_until_truthy(timeout:, interval: 1)
    start_time = Process.clock_gettime(Process::CLOCK_MONOTONIC, :second)
    ending_time = start_time + timeout

    loop do
      result = yield
      return result if result

      remaining = ending_time - Process.clock_gettime(Process::CLOCK_MONOTONIC, :second)
      break if remaining <= 0

      sleep [interval, remaining].min
    end

    nil
  end
end
