# -*- 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
  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:

protected



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

    begin
      exploit, payload = ctx

      # 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

      # Build a user-friendly error message
      msg = "#{e}"
      unless e.class == Msf::Exploit::Failed
        msg = "#{e.class} #{e}"
      end

      exploit.error = e

      # Record the detailed reason
      exploit.fail_detail ||= e.to_s

      case e
      when Msf::Exploit::Complete
        # Nothing to show in this case
        return

      when Msf::Exploit::Failed
        exploit.print_error("Exploit aborted due to failure: #{exploit.fail_reason}: #{msg}")

        # The caller should have already set exploit.fail_reason
        if exploit.fail_reason == Msf::Exploit::Failure::None
          exploit.fail_reason = Msf::Exploit::Failure::Unknown
        end

      when Rex::ConnectionError
        exploit.fail_reason = Msf::Exploit::Failure::Unreachable
        exploit.print_error("Exploit failed [#{exploit.fail_reason}]: #{msg}")
        elog("Exploit failed (#{exploit.refname}): #{msg}", 'core', LEV_0)
        dlog("Call stack:\n#{e.backtrace.join("\n")}", 'core', LEV_3)

      when Timeout::Error
        exploit.fail_reason = Msf::Exploit::Failure::TimeoutExpired
        exploit.print_error("Exploit failed [#{exploit.fail_reason}]: #{msg}")
        elog("Exploit failed (#{exploit.refname}): #{msg}", 'core', LEV_0)
        dlog("Call stack:\n#{e.backtrace.join("\n")}", 'core', LEV_3)
      else

        # Compare as a string since not all error classes may be loaded
        case msg
        when /access.denied|Login Failed/i # Covers SMB as well as some generic errors
          exploit.fail_reason = Msf::Exploit::Failure::NoAccess
        when /connection reset/i
          exploit.fail_reason = Msf::Exploit::Failure::Disconnected
        when /connection timed out|SSL_connect|unreachable|connection was refused/i
          exploit.fail_reason = Msf::Exploit::Failure::Unreachable
        when /unable.*target/i
          exploit.fail_reason = Msf::Exploit::Failure::NoTarget
        when /execution expired/i
          exploit.fail_reason = Msf::Exploit::Failure::TimeoutExpired
        when /(doesn.t|not).*vulnerable|may.*patched/i
          exploit.fail_reason = Msf::Exploit::Failure::NotVulnerable
        end

        # The caller should have already set exploit.fail_reason
        if exploit.fail_reason == Msf::Exploit::Failure::None
          exploit.fail_reason = Msf::Exploit::Failure::Unknown
        end

        if exploit.fail_reason == Msf::Exploit::Failure::Unknown
          exploit.print_error("Exploit failed: #{msg}")
        else
          exploit.print_error("Exploit failed [#{exploit.fail_reason}]: #{msg}")
        end

        elog("Exploit failed (#{exploit.refname}): #{msg}", 'core', LEV_0)
        dlog("Call stack:\n#{e.backtrace.join("\n")}", 'core', LEV_3)
      end

      # Record the error to various places
      exploit.framework.events.on_module_error(exploit, msg)

      # Report the failure (and attempt) in the database
      exploit.report_failure
    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
    payload.stop_handler

    exploit.framework.events.on_module_complete(exploit)

    # Allow the exploit to cleanup after itself, that messy bugger.
    exploit.cleanup
  end

end

end