# -*- coding: binary -*-

#
# This mixin implements an exploit's check method by invoking an aux module
#

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 it isn't aux
    if mod.type != Msf::MODULE_AUX
      return CheckCode::Unsupported(
        "#{check_module} is not an auxiliary module"
      )
    end

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

    print_status("Using #{check_module} as check")

    # Retrieve the module's return value
    res = mod.run_simple(
      'LocalInput'  => user_input,
      'LocalOutput' => user_output,
      'Options'     => datastore # XXX: This clobbers the datastore!
    )

    # Ensure return value is a CheckCode
    case res
    when Exploit::CheckCode
      # Return the CheckCode
      res
    when Hash
      # XXX: Find CheckCode associated with RHOST, which is set automatically
      checkcode = res[datastore['RHOST']]

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

      # Return the CheckCode
      checkcode
    else
      # Bail if module doesn't return a CheckCode
      Exploit::CheckCode::Unsupported(
        "#{check_module} does not return a CheckCode"
      )
    end
  end

  def check_module
    datastore['CheckModule']
  end

end
end