Persistence Exploit Suggester

##
# This module requires Metasploit: https://metasploit.com/download
# Current source: https://github.com/rapid7/metasploit-framework
##

class MetasploitModule < Msf::Post

  include Msf::Auxiliary::Report

  def initialize(info = {})
    super(
      update_info(
        info,
        'Name' => 'Persistence Exploit Suggester',
        'Description' => %q{
          This module suggests persistence modules that can be used.
          The modules are suggested based on the architecture and platform
          that the user has a shell opened as well as the available exploits
          in meterpreter.
          It's important to note that not all modules will be checked.
          Exploits are chosen based on these conditions: session type,
          platform, architecture, and required default options.
        },
        'License' => MSF_LICENSE,
        'Author' => [ 'h00die' ],
        'Platform' => all_platforms,
        'SessionTypes' => [ 'meterpreter', 'shell' ],
        'Notes' => {
          'Stability' => [],
          'Reliability' => [],
          'SideEffects' => []
        }
      )
    )
    register_options [
      Msf::OptInt.new('SESSION', [ true, 'The session to run this module on' ]),
      Msf::OptBool.new('SHOWDESCRIPTION', [true, 'Displays a detailed description for the available exploits', false])
    ]

    register_advanced_options(
      [
        # most linux persistence modules are arch-cmd but for payload purposes only
        # but usually we end up with a meterpreter session, thus making these invalid
        # so disable this check by default
        Msf::OptBool.new('ValidateArch', [true, 'Validate architecture', false]),
        Msf::OptBool.new('ValidatePlatform', [true, 'Validate platform', true]),
        Msf::OptBool.new('ValidateMeterpreterCommands', [true, 'Validate Meterpreter commands', false]),
        # https://github.com/rapid7/rex-text/blob/a72151d409cd812978f63ad0c330efbc8f44b977/lib/rex/text/color.rb#L13
        Msf::OptString.new('Colors', [false, 'Valid, Invalid and Ignored colors for module checks (unset to disable)', 'grn/red/blu'])
      ]
    )
  end

  def valid_colors?(color_str = datastore['Colors'])
    tokens = color_str.split('/')
    tokens.each do |tok|
      print_warning "#{tok} is unlikely to have any functionality for printing colors." if tok == 'clr'

      unless Rex::Text::Color::SUPPORTED_FORMAT_CODES.include?("%#{tok}")
        print_error "#{tok} is NOT valid color. Please see https://github.com/rapid7/rex-text/blob/a72151d409cd812978f63ad0c330efbc8f44b977/lib/rex/text/color.rb#L13 for valid color options"
        return false
      end
    end
    true
  end

  def all_platforms
    Msf::Module::Platform.subclasses.collect { |c| c.realname.downcase }
  end

  def session_arch
    # Prefer calling native arch when available, as most LPEs will require this (e.g. x86, x64) as opposed to Java/Python Meterpreter's values (e.g. Java, Python)
    session.respond_to?(:native_arch) ? session.native_arch : session.arch
  end

  def is_module_arch?(mod)
    mod_arch = mod.target.arch || mod.arch
    mod_arch.include?(session_arch)
  rescue StandardError => e
    print_error "Failed to check module arch for #{mod.fullname} => #{e}"
  end

  def is_module_wanted?(mod)
    mod[:result][:incompatibility_reasons].empty?
  end

  def is_session_type?(mod)
    # There are some modules that do not define any compatible session types.
    # We could assume that means the module can run on all session types,
    # Or we could consider that as incorrect module metadata.
    mod.session_types.include?(session.type)
  end

  def is_module_platform?(mod)
    return false if mod.target.nil?

    platform_obj = Msf::Module::Platform.find_platform session.platform

    module_platforms = mod.target.platform ? mod.target.platform.platforms : mod.platform.platforms
    module_platforms.include? platform_obj
  rescue ArgumentError => e
    # When not found, find_platform raises an ArgumentError
    elog('Could not find a platform', error: e)
    return false
  end

  def has_required_module_options?(mod)
    get_all_missing_module_options(mod).empty?
  end

  def get_all_missing_module_options(mod)
    missing_options = []
    mod.options.each_pair do |option_name, option|
      missing_options << option_name if option.required && option.default.nil? && mod.datastore[option_name].blank?
    end
    missing_options
  end

  def valid_incompatibility_reasons(mod, verify_reasons)
    # As we can potentially ignore some `reasons` (e.g. accepting arch values which are, on paper, not compatible),
    # this keeps track of valid reasons why we will not consider the module that we are evaluating to be valid.
    valid_reasons = []
    valid_reasons << "Missing required module options (#{get_all_missing_module_options(mod).join('. ')})" unless verify_reasons[:has_required_module_options]

    incompatible_opts = []
    incompatible_opts << 'architecture' unless verify_reasons[:is_module_arch]
    incompatible_opts << 'platform' unless verify_reasons[:is_module_platform]
    incompatible_opts << 'session type' unless verify_reasons[:is_session_type]
    valid_reasons << "Not Compatible (#{incompatible_opts.join(', ')})" if incompatible_opts.any?

    valid_reasons << 'Missing/unloadable Meterpreter commands' if verify_reasons[:missing_meterpreter_commands].any?
    valid_reasons
  end

  def set_module_options(mod)
    ignore_list = ['ACTION', 'TARGET'].freeze
    datastore.each_pair do |k, v|
      mod.datastore[k] = v unless ignore_list.include?(k.upcase)
    end
    if !mod.datastore['SESSION'] && session.present?
      mod.datastore['SESSION'] = session.sid
    end
  end

  def set_module_target(mod)
    session_platform = Msf::Module::Platform.find_platform(session.platform)
    target_index = mod.targets.find_index do |target|
      # If the target doesn't define its own compatible platforms or architectures, default to the parent (module) values.
      target_platforms = target.platform&.platforms || mod.platform.platforms
      target_architectures = target.arch || mod.arch

      target_platforms.include?(session_platform) && target_architectures.include?(session_arch)
    end
    mod.datastore['Target'] = target_index if target_index
  end

  def setup
    return unless session

    print_status "Collecting persistence modules for #{session.session_type}..."

    setup_validation_options
    if valid_colors?
      setup_color_options
    else
      fail_with(Failure::BadConfig, 'Colors options set incorrectly')
    end

    # Collects persistence modules into an array
    @persistence_modules = []
    exploit_refnames = framework.exploits.module_refnames
    exploit_refnames.each_with_index do |name, index|
      print "%bld%blu[*]%clr Collecting exploit #{index + 1} / #{exploit_refnames.count}\r"
      next unless name.include? '/persistence/'

      mod = framework.exploits.create name
      next unless mod

      set_module_options mod
      set_module_target mod
      verify_result = verify_mod(mod)
      @persistence_modules << { module: mod, result: verify_result } if verify_result[:has_check]
    end
  end

  def verify_mod(mod)
    return { has_check: false } unless mod.is_a?(Msf::Exploit::Local) && mod.has_check?

    result = {
      has_check: true,
      is_module_platform: (@validate_platform ? is_module_platform?(mod) : true),
      is_module_arch: (@validate_arch ? is_module_arch?(mod) : true),
      has_required_module_options: has_required_module_options?(mod),
      missing_meterpreter_commands: (@validate_meterpreter_commands && session.type == 'meterpreter') ? meterpreter_session_incompatibility_reasons(session) : [],
      is_session_type: is_session_type?(mod)
    }
    result[:incompatibility_reasons] = valid_incompatibility_reasons(mod, result)
    result
  end

  def setup_validation_options
    @validate_arch = datastore['ValidateArch']
    @validate_platform = datastore['ValidatePlatform']
    @validate_meterpreter_commands = datastore['ValidateMeterpreterCommands']
  end

  def setup_color_options
    @valid_color, @invalid_color, @ignored_color =
      (datastore['Colors'] || '').split('/')

    @valid_color = "%#{@valid_color}" unless @valid_color.blank?
    @invalid_color = "%#{@invalid_color}" unless @invalid_color.blank?
    @ignored_color = "%#{@ignored_color}" unless @ignored_color.blank?
  end

  def show_found_exploits
    unless datastore['VERBOSE']
      print_status "#{@persistence_modules.length} exploit checks are being tried..."
      return
    end

    vprint_status "The following #{@persistence_modules.length} exploit checks are being tried:"
    @persistence_modules.each do |x|
      vprint_status x[:module].fullname
    end
  end

  def run
    runnable_exploits = @persistence_modules.select { |mod| is_module_wanted?(mod) }
    if runnable_exploits.empty?
      print_error 'No suggestions available.'
      vprint_line
      vprint_session_info
      vprint_status unwanted_modules_table(@persistence_modules.reject { |mod| is_module_wanted?(mod) })
      return
    end

    show_found_exploits
    results = runnable_exploits.map.with_index do |mod, index|
      print "%bld%blu[*]%clr Running check method for exploit #{index + 1} / #{runnable_exploits.count}\r"
      begin
        checkcode = mod[:module].check_simple
      rescue StandardError => e
        elog("#Local Persistence Suggester failed with: #{e.class} when using #{mod[:module].shortname}", error: e)
        vprint_error "Check with module #{mod[:module].fullname} failed with error #{e.class}"
        next { module: mod[:module], errors: ['The check raised an exception.'] }
      end

      if checkcode.nil?
        vprint_error "Check failed with #{mod[:module].fullname} for unknown reasons"
        next { module: mod[:module], errors: ['The check failed for unknown reasons.'] }
      end

      # See def is_check_interesting?
      unless is_check_interesting? checkcode
        vprint_status "#{mod[:module].fullname}: #{checkcode.message}"
        next { module: mod[:module], errors: [checkcode.message] }
      end

      # Prints the full name and the checkcode message for the exploit
      print_good "#{mod[:module].fullname}: #{checkcode.message}"

      # If the datastore option is true, a detailed description will show
      if datastore['SHOWDESCRIPTION']
        # Formatting for the description text
        Rex::Text.wordwrap(Rex::Text.compress(mod[:module].description), 2, 70).split(/\n/).each do |line|
          print_line line
        end
      end

      next { module: mod[:module], checkcode: checkcode.message }
    end

    print_line
    print_status valid_modules_table(results)

    vprint_line
    vprint_session_info
    vprint_status unwanted_modules_table(@persistence_modules.reject { |mod| is_module_wanted?(mod) })

    report_data = {}
    results.each do |result|
      report_data[result[:module].fullname] = result[:checkcode] if result[:checkcode]
    end

    report_note({
      host: session.session_host,
      type: 'persistence.suggested_module',
      data: report_data
    })
  end

  def valid_modules_table(results)
    name_styler = ::Msf::Ui::Console::TablePrint::CustomColorStyler.new
    check_styler = ::Msf::Ui::Console::TablePrint::CustomColorStyler.new

    # Split all the results by their checkcode.
    # We want the modules that returned a checkcode to be at the top.
    checkcode_rows, without_checkcode_rows = results.partition { |result| result[:checkcode] }
    rows = (checkcode_rows + without_checkcode_rows).map.with_index do |result, index|
      color = result[:checkcode] ? @valid_color : @invalid_color
      check_res = result.fetch(:checkcode) { result[:errors].join('. ') }
      name_styler.merge!({ result[:module].fullname => color })
      check_styler.merge!({ check_res => color })

      [
        index + 1,
        result[:module].fullname,
        result[:checkcode] ? 'Yes' : 'No',
        check_res
      ]
    end

    Rex::Text::Table.new(
      'Header' => "Valid modules for session #{session.sid}:",
      'Indent' => 1,
      'Columns' => [ '#', 'Name', 'Potentially Vulnerable?', 'Check Result' ],
      'SortIndex' => -1,
      'WordWrap' => false, # Don't wordwrap as it messes up coloured output when it is broken up into more than one line
      'ColProps' => {
        'Name' => {
          'Stylers' => [name_styler]
        },
        'Potentially Vulnerable?' => {
          'Stylers' => [::Msf::Ui::Console::TablePrint::CustomColorStyler.new({ 'Yes' => @valid_color, 'No' => @invalid_color })]
        },
        'Check Result' => {
          'Stylers' => [check_styler]
        }
      },
      'Rows' => rows
    )
  end

  def unwanted_modules_table(unwanted_modules)
    arch_styler = ::Msf::Ui::Console::TablePrint::CustomColorStyler.new
    platform_styler = ::Msf::Ui::Console::TablePrint::CustomColorStyler.new
    session_type_styler = ::Msf::Ui::Console::TablePrint::CustomColorStyler.new

    rows = unwanted_modules.map.with_index do |mod, index|
      begin
        platforms = mod[:module].target.platform&.platforms&.any? ? mod[:module].target.platform.platforms : mod[:module].platform.platforms
      rescue NoMethodError
        platforms = nil
      end
      platforms ||= []
      begin
        arch = mod[:module].target.arch&.any? ? mod[:module].target.arch : mod[:module].arch
      rescue NoMethodError
        arch = nil
      end
      arch ||= []

      arch.each do |a|
        if a != session_arch
          if @validate_arch
            color = @invalid_color
          else
            color = @ignored_color
          end
        else
          color = @valid_color
        end

        arch_styler.merge!({ a.to_s => color })
      end

      platforms.each do |module_platform|
        if module_platform != ::Msf::Module::Platform.find_platform(session.platform)
          if @validate_platform
            color = @invalid_color
          else
            color = @ignored_color
          end
        else
          color = @valid_color
        end

        platform_styler.merge!({ module_platform.realname => color })
      end

      mod[:module].session_types.each do |session_type|
        color = session_type == session.type ? @valid_color : @invalid_color
        session_type_styler.merge!(session_type.to_s => color)
      end

      [
        index + 1,
        mod[:module].fullname,
        mod[:result][:incompatibility_reasons].join('. '),
        platforms.any? ? platforms.map(&:realname).sort.join(', ') : 'No defined platforms',
        arch.any? ? arch.sort.join(', ') : 'No defined architectures',
        mod[:module].session_types.any? ? mod[:module].session_types.sort.join(', ') : 'No defined session types'
      ]
    end

    Rex::Text::Table.new(
      'Header' => "Incompatible modules for session #{session.sid}:",
      'Indent' => 1,
      'Columns' => [ '#', 'Name', 'Reasons', 'Platform', 'Architecture', 'Session Type' ],
      'WordWrap' => false,
      'ColProps' => {
        'Architecture' => {
          'Stylers' => [arch_styler]
        },
        'Platform' => {
          'Stylers' => [platform_styler]
        },
        'Session Type' => {
          'Stylers' => [session_type_styler]
        }
      },
      'Rows' => rows
    )
  end

  def vprint_session_info
    vprint_status 'Current Session Info:'
    vprint_status "  Session Type: #{session.type}"
    vprint_status "  Architecture: #{session_arch}"
    vprint_status "  Platform: #{session.platform}"
  end

  def is_check_interesting?(checkcode)
    [
      Msf::Exploit::CheckCode::Vulnerable,
      Msf::Exploit::CheckCode::Appears,
      Msf::Exploit::CheckCode::Detected
    ].include? checkcode
  end

  def print_status(msg = '')
    super(session ? "#{session.session_host} - #{msg}" : msg)
  end

  def print_good(msg = '')
    super(session ? "#{session.session_host} - #{msg}" : msg)
  end

  def print_error(msg = '')
    super(session ? "#{session.session_host} - #{msg}" : msg)
  end
end