# -*- coding: binary -*-

module Msf
module Simple

###
#
# A simplified exploit wrapper.
#
###
module Exploit

  include Module

  #
  # Wraps the exploitation process in a simple single method.  The options
  # hash can have the following values passed in it:
  #
  # Encoder
  #
  # 	The encoder module that should be used.
  #
  # Payload
  #
  # 	The payload module name that should be used.
  #
  # Target
  #
  # 	The selected target index.
  #
  # Nop
  #
  # 	The NOP generator that should be used in preference.
  #
  # OptionStr
  #
  # 	A string of comma separated option values that should be imported into
  # 	the datastore.
  #
  # Options
  #
  # 	A hash of values to be imported directly into the datastore.
  #
  # LocalInput
  #
  # 	The local input handle that data can be read in from.
  #
  # LocalOutput
  #
  # 	The local output through which data can be displayed.
  #
  # RunAsJob
  #
  # 	Whether or not the exploit should be run in the context of a background
  # 	job.
  #
  def self.exploit_simple(oexploit, opts, &block)
    exploit = oexploit.replicant
    # Trap and print errors here (makes them UI-independent)
    begin
      # Clone the module to prevent changes to the original instance

      Msf::Simple::Framework.simplify_module(exploit)
      yield(exploit) if block_given?

      # Import options from the OptionStr or Option hash.
      exploit._import_extra_options(opts)
      opts['Payload'] ||= exploit.datastore['Payload']

      unless opts['Quiet']
        exploit.init_ui(opts['LocalInput'] || exploit.user_input, opts['LocalOutput'] || exploit.user_output)
      else
        exploit.init_ui(nil, nil)
      end

      # Make sure parameters are valid.
      if (opts['Payload'] == nil)
        raise MissingPayloadError, 'A payload has not been selected.', caller
      end

      # Verify the options
      exploit.options.validate(exploit.datastore)

      # Start it up
      driver = Msf::ExploitDriver.new(exploit.framework)

      # Keep the handler of driver running if exploiting multiple targets.
      driver.keep_handler = true if opts['multi']

      # Initialize the driver instance
      driver.exploit    = exploit
      driver.payload    = exploit.framework.payloads.create(opts['Payload'])

      # Set the force wait for session flag if the caller requested force
      # blocking.  This is so that passive exploits can be blocked on from
      # things like the cli.
      driver.force_wait_for_session = true if (opts['ForceBlocking'] == true)

      # Was the payload valid?
      if (driver.payload == nil)
        raise MissingPayloadError,
          "You specified an invalid payload: #{opts['Payload']}", caller
      end

      # Use the supplied encoder, if any.  If one was not specified, then
      # nil will be assigned causing the exploit to default to picking the
      # best encoder.
      exploit.datastore['ENCODER'] = opts['Encoder'] if opts['Encoder']

      # Use the supplied NOP generator, if any.  If one was not specified, then
      # nil will be assigned causing the exploit to default to picking a
      # compatible NOP generator.
      exploit.datastore['NOP'] = opts['Nop'] if opts['Nop']

      # Force the payload to share the exploit's datastore
      driver.payload.share_datastore(driver.exploit.datastore)

      # Verify the payload options
      driver.payload.options.validate(driver.payload.datastore)

      # Set the target and then work some magic to derive index
      exploit.datastore['TARGET'] = opts['Target'] if opts['Target']
      target_idx = exploit.target_index

      if (target_idx == nil or target_idx < 0)
        raise MissingTargetError,
          "You must select a target.", caller
      end

      driver.target_idx = target_idx

      # Set the payload and exploit's subscriber values
      unless opts['Quiet']
        driver.payload.init_ui(opts['LocalInput'] || exploit.user_input, opts['LocalOutput'] || exploit.user_output)
      else
        driver.payload.init_ui(nil, nil)
      end

      if (opts['RunAsJob'])
        driver.use_job = true
      end

      # Let's rock this party
      driver.run

      # Save the job identifier this exploit is running as
      exploit.job_id  = driver.job_id

      # Propagate this back to the caller for console mgmt
      oexploit.job_id = exploit.job_id
    rescue ::Interrupt
      exploit.error = $!
      raise $!
    rescue ::Msf::OptionValidateError => e
      exploit.error = e
      ::Msf::Ui::Formatter::OptionValidateError.print_error(exploit, e)
      return false
    rescue ::Exception => e
      exploit.error = e
      exploit.print_error("Exploit failed: #{e}")
      elog("Exploit failed (#{exploit.refname})", error: e)
    end

    return driver.session if driver
    nil
  end

  #
  # Calls the class method.
  #
  def exploit_simple(opts, &block)
    Msf::Simple::Exploit.exploit_simple(self, opts, &block)
  end

  alias run_simple exploit_simple
  #
  # Initiates a check, setting up the exploit to be used.  The following
  # options can be specified:
  #
  # LocalInput
  #
  # 	The local input handle that data can be read in from.
  #
  # LocalOutput
  #
  # 	The local output through which data can be displayed.
  #
  def self.check_simple(mod, opts, job_listener: Msf::Simple::NoopJobListener.instance)
    Msf::Simple::Framework.simplify_module(mod)
    mod._import_extra_options(opts)

    if opts['LocalInput']
      mod.init_ui(opts['LocalInput'], opts['LocalOutput'])
    end

    unless mod.has_check?
      # Bail out early if the module doesn't have check
      raise ::NotImplementedError.new(Msf::Exploit::CheckCode::Unsupported.message)
    end

    # Validate the option container state so that options will
    # be normalized
    mod.validate

    run_uuid = Rex::Text.rand_text_alphanumeric(24)
    job_listener.waiting run_uuid
    ctx = [mod, run_uuid, job_listener]

    if opts['RunAsJob']
      mod.job_id = mod.framework.jobs.start_bg_job(
        "Exploit: #{mod.refname} check",
        ctx,
        Proc.new { |ctx_| self.job_check_proc(ctx_) },
        Proc.new { |ctx_| nil }
      )
      [run_uuid, mod.job_id]
    else
      self.job_check_proc(ctx)
    end
  end

  #
  # Calls the class method.
  #
  def check_simple(opts)
    Msf::Simple::Exploit.check_simple(self, opts)
  end

  protected

  def self.job_check_proc(ctx)
    mod = ctx[0]
    run_uuid = ctx[1]
    job_listener = ctx[2]
    begin
      job_listener.start run_uuid
      mod.setup
      result = mod.check
      job_listener.completed(run_uuid, result, mod)
    rescue => e
      job_listener.failed(run_uuid, e, mod)
      mod.handle_exception e
    ensure
      mod.cleanup
    end

    return result
  end
end

end
end