# -*- coding: binary -*-

#
# This mixin implements an exploit's check method by invoking an aux module
#
# NOTE: The module's run_host/run method MUST return an Msf::Exploit::CheckCode
#

module Msf
module Exploit::Remote::CheckModule

  def initialize(info = {})
    super

    register_advanced_options([
      OptString.new('CheckModule', [true, 'Module to check with'])
    ])
  end

  # When this mixin is included, this method becomes the exploit's check method
  #
  # @return [Msf::Exploit::CheckCode] Whether or not the target is vulnerable
  def check
    # Instantiate the module
    mod = framework.modules.create(check_module)

    # Bail if we couldn't
    unless mod
      return CheckCode::Unsupported("Could not instantiate #{check_module}")
    end

    # Bail if run_host/run isn't defined
    if mod.respond_to?(:run_host)
      meth = :run_host
    elsif mod.respond_to?(:run)
      meth = :run
    else
      return CheckCode::Unsupported("#{check_module} does not define a run_host/run method")
    end

    # Add the exploit's targeting options to the module's datastore
    %w[RHOSTS RHOST RPORT].each do |opt|
      next unless datastore[opt]

      mod.datastore[opt] = datastore[opt].dup
    end

    # Bail if module options don't validate
    mod.options.validate(mod.datastore)

    # Use the exploit's input and output as the module's
    mod.user_input, mod.user_output = user_input, user_output

    # Use the module's CheckCode
    checkcode =
      case meth
      when :run_host
        mod.run_host(rhost)
      when :run
        mod.run
      end

    # Bail if module doesn't return a CheckCode
    unless checkcode.kind_of?(Exploit::CheckCode)
      print_warning("#{check_module} does not return a CheckCode")
      return Exploit::CheckCode::Unsupported
    end

    # Return the CheckCode
    checkcode
  end

  def check_module
    datastore['CheckModule']
  end

end
end