module Msf
  module Exploit::AutoTarget

    # Checks to see if the auto-generated Automatic Targeting
    # has been selected. If the module had an already defined
    # Automatic target, then we let the module handle the targeting
    # itself.
    #
    # @return [Boolean] whether or not to use our automatic targeting routine
    def auto_target?
      selected_target = targets[target_index]
      return false if selected_target.nil?
      if selected_target.name =~ /Automatic/ && selected_target['AutoGenerated'] == true  && auto_target_host
        true
      else
        false
      end
    end

    # Returns the Target Index of the automatically selected Target from
    # our Automatic Targeting routine.
    #
    # @return [Integer] the index of the selected Target
    # @return [nil] if no target could be selected
    def auto_targeted_index
      selected_target = select_target
      return nil if selected_target.nil?
      targets.each_with_index do |target, index|
        return index if target == selected_target
      end
      nil
    end

    # Chooses the best possible Target for what we know about
    # the targeted host.
    #
    # @return [Msf::Module::Target] the Target that our automatic routine selected
    def select_target
      return nil unless auto_target?
      host_record = auto_target_host
      return nil if host_record.nil?
      filtered_targets = filter_by_os(host_record)
      filtered_targets.first
    end

    # Finds an <Mdm::Host> for the RHOST if one exists
    #
    # @return [Mdm:Host] the Host record if one exists
    # @return [nil] if no Host record is present, or the DB is not active
    def auto_target_host
      return nil unless self.respond_to?(:rhost)
      return nil unless framework.db.active
      current_workspace = framework.db.find_workspace(self.workspace)
      current_workspace.hosts.where(address: rhost).first
    end

    # Returns the best matching Targets based on the target host's
    # OS information. It looks at the OS Family, OS Name, and OS SP.
    #
    # @param host_record [Mdm::Host] the target host record
    # @return [Array<Msf::Module::Target>] an array of matching targets
    def filter_by_os(host_record)
      filtered_by_family = filter_by_os_family(host_record)
      filtered_by_name   = filter_by_os_name(filtered_by_family, host_record)
      # If Filtering by name gave us no results, then we reset back to the family filter group
      filtered_by_name   = filtered_by_family if filtered_by_name.empty?
      filtered_by_sp     = filter_by_os_sp(filtered_by_name,host_record)
      # If Filtering by SP was a bust, revert back one level
      filtered_by_sp     = filtered_by_name if filtered_by_sp.empty?
      filtered_by_sp
    end

    # Returns all Targets that match the target host's OS Family
    # e.g Windows, Linux, OS X, etc
    #
    # @param host_record [Mdm::Host] the target host record
    # @return [Array<Msf::Module::Target>] an array of matching targets
    def filter_by_os_family(host_record)
      return [] if host_record.os_family.blank?
      filtered_targets = targets.collect do |target|
        if target.name =~ /#{host_record.os_family}/
          target
        else
          nil
        end
      end
      filtered_targets.compact
    end

    # Returns all Targets that match the target host's OS Name
    # e.g Windows 7, Windows XP, Windows Vista, etc
    #
    # @param potential_targets [Array<Msf::Module::Target>] the filtered targets that we wish to filter further
    # @param host_record [Mdm::Host] the target host record
    # @return [Array<Msf::Module::Target>] an array of matching targets
    def filter_by_os_name(potential_targets, host_record)
      return [] if host_record.os_name.blank?
      filtered_targets = []
      potential_targets.each do |target|
        filtered_targets << target if target.name =~ /#{host_record.os_name}/
      end
      filtered_targets
    end

    # Returns all Targets that match the target host's OS SP
    #
    # @param potential_targets [Array<Msf::Module::Target>] the filtered targets that we wish to filter further
    # @param host_record [Mdm::Host] the target host record
    # @return [Array<Msf::Module::Target>] an array of matching targets
    def filter_by_os_sp(potential_targets, host_record)
      return [] if host_record.os_sp.blank?
      filtered_targets = []
      potential_targets.each do |target|
        filtered_targets << target if target.name =~ /#{host_record.os_sp}/
      end
      filtered_targets
    end
  end
end