# -*- coding: binary -*-
require 'msf/core'

module Msf

###
#
# This class drives the exploitation process from start to finish for a given
# exploit module instance.  It's responsible for payload generation, encoding,
# and padding as well as initialization handlers and finally launching the
# exploit.
#
###
class ExploitDriver

  #
  # Initializes the exploit driver using the supplied framework instance.
  #
  def initialize(framework)
    self.payload                = nil
    self.exploit                = nil
    self.use_job                = false
    self.job_id                 = nil
    self.force_wait_for_session = false
    self.semaphore              = Mutex.new
  end

  #
  # Specification of the exploit target index.
  #
  def target_idx=(target_idx)
    if (target_idx)
      # Make sure the target index is valid
      if (target_idx >= exploit.targets.length)
        raise Rex::ArgumentError, "Invalid target index.", caller
      end
    end

     # Set the active target
    @target_idx = target_idx
  end

  #
  # This method returns the currently selected target index.
  #
  def target_idx
    @target_idx
  end

  #
  # Checks to see if the supplied payload is compatible with the
  # current exploit.  Assumes that target_idx is valid.
  #
  def compatible_payload?(payload)
    # Try to use the target's platform in preference of the exploit's
    exp_platform = exploit.targets[target_idx].platform || exploit.platform

    return ((payload.platform & exp_platform).empty? == false)
  end

  ##
  #
  # Exploit execution
  #
  ##

  #
  # Makes sure everything's in tip-top condition prior to launching the
  # exploit.  For things that aren't good to go, an exception is thrown.
  #
  def validate
    # First, validate that a target has been selected
    if (target_idx == nil)
      raise MissingTargetError,
        "A payload cannot be selected until a target is specified.",
        caller
    end

    # Next, validate that a payload has been selected
    if (payload == nil)
      raise MissingPayloadError,
        "A payload has not been selected.", caller
    end

    # Make sure the payload is compatible after all
    if (compatible_payload?(payload) == false)
      raise IncompatiblePayloadError.new(payload.refname),
        "Incompatible payload", caller
    end

    # Associate the payload instance with the exploit
    payload.assoc_exploit = exploit

    # Finally, validate options on the exploit module to ensure that things
    # are ready to operate as they should.
    exploit.options.validate(exploit.datastore)

    # Validate the payload's options.  The payload's datastore is
    # most likely shared against the exploit's datastore, but in case it
    # isn't.
    payload.options.validate(payload.datastore)

    return true
  end

  #
  # Kicks off an exploitation attempt and performs the following four major
  # operations:
  #
  #   - Generates the payload
  #   - Initializes & monitors the handler
  #   - Launches the exploit
  #   - Cleans up the handler
  #
  def run
    # First thing's first -- validate the state.  Make sure all requirement
    # parameters are set, including those that are derived from the
    # datastore.
    validate()

    # After validation has occurred, it's time to set some values on the
    # exploit instance and begin preparing the payload
    exploit.datastore['TARGET'] = target_idx

    # Default the session to nil
    self.session = nil

    # Explicitly clear the module's job_id in case it was set in a previous
    # run
    exploit.job_id = nil

    # If we are being instructed to run as a job then let's create that job
    # like a good person.
    if (use_job or exploit.stance == Msf::Exploit::Stance::Passive)
      # Since references to the exploit and payload will hang around for
      # awhile in the job, make sure we copy them so further changes to
      # the datastore don't alter settings in existing jobs
      e = exploit.replicant
      p = payload.replicant

      # Assign the correct exploit instance to the payload
      p.assoc_exploit = e

      # Generate the encoded version of the supplied payload for the
      # newly copied exploit module instance
      e.generate_payload(p)
      ctx = [ e, p ]

      e.job_id = e.framework.jobs.start_bg_job(
        "Exploit: #{e.refname}",
        ctx,
        Proc.new { |ctx_| job_run_proc(ctx_) },
        Proc.new { |ctx_| job_cleanup_proc(ctx_) }
      )
      self.job_id = e.job_id
    else
      # Generate the encoded version of the supplied payload on the
      # exploit module instance
      exploit.generate_payload(payload)

      # No need to copy since we aren't creating a job.  We wait until
      # they're finished running to do anything else with them, so
      # nothing should be able to modify their datastore or other
      # settings until after they're done.
      ctx = [ exploit, payload ]

      job_run_proc(ctx)
      job_cleanup_proc(ctx)
    end

    return session
  end

  attr_accessor :exploit # :nodoc:
  attr_accessor :payload # :nodoc:
  attr_accessor :use_job # :nodoc:
  #
  # The identifier of the job this exploit is launched as, if it's run as a
  # job.
  #
  attr_accessor :job_id
  attr_accessor :force_wait_for_session # :nodoc:
  attr_accessor :session # :nodoc:

  # To synchronize threads cleaning up the exploit and the handler
  attr_accessor :semaphore

protected

  #
  # Job run proc, sets up the exploit and kicks it off.
  #
  def job_run_proc(ctx)
    begin
      exploit, payload = ctx
      # Default session wait time..
      delay = payload.wfs_delay + exploit.wfs_delay
      delay = nil if exploit.passive?

      # Set the exploit up the bomb
      exploit.setup

      exploit.framework.events.on_module_run(exploit)

      # Launch the exploit
      exploit.exploit

    rescue ::Exception => e
      if [::RuntimeError, ::Interrupt].include?(e.class)
        # Wait for session, but don't wait long.
        delay = 0.01
      end

      exploit.handle_exception e
    end

    # Wait the payload to acquire a session if this isn't a passive-style
    # exploit.
    return if not delay


    if (force_wait_for_session == true) or
      (exploit.passive? == false and exploit.handler_enabled?)
      begin
        self.session = payload.wait_for_session(delay)
      rescue ::Interrupt
        # Don't let interrupt pass upward
      end
    end

    return self.session if self.session

    if exploit.fail_reason == Msf::Exploit::Failure::None
      exploit.fail_reason = Msf::Exploit::Failure::PayloadFailed
      exploit.fail_detail = "No session created"
      exploit.report_failure
    end
  end

  #
  # Clean up the exploit and the handler after the job completes.
  #
  def job_cleanup_proc(ctx)
    exploit, payload = ctx

    # Ensure that, no matter what, clean up of the handler occurs
    semaphore.synchronize { payload.stop_handler }

    exploit.framework.events.on_module_complete(exploit)

    # Allow the exploit to cleanup after itself, that messy bugger.
    semaphore.synchronize { exploit.cleanup }
  end

end

end