# frozen_string_literal: true

module Msf
  module SessionCompatibility
    include Msf::Auxiliary::Report

    include Msf::Module::HasActions
    include Msf::Post::Common

    def initialize(info = {})
      super

      # Default stance is active
      self.passive = info['Passive'] || false
      self.session_types = info['SessionTypes'] || []
    end

    #
    # Grabs a session object from the framework or raises {OptionValidateError}
    # if one doesn't exist and is required. Initializes user input and output on the session.
    #
    # @raise [OptionValidateError] if {#session} returns nil
    def setup
      alert_user

      if options['SESSION']&.required && session.blank?
        raise Msf::OptionValidateError, ['SESSION']
      end

      if datastore['SESSION'] && session.nil?
        raise Msf::OptionValidateError, ['SESSION']
      end

      # Msf::Exploit#setup for exploits, NoMethodError for post modules
      super rescue NoMethodError

      return unless session

      # Check session readiness before compatibility so the session can be queried
      # for its platform, capabilities, etc.
      check_for_session_readiness if session.type == "meterpreter"

      incompatibility_reasons = session_incompatibility_reasons(session)
      if incompatibility_reasons.any?
        print_warning('SESSION may not be compatible with this module:')
        incompatibility_reasons.each do |reason|
          print_warning(" * #{reason}")
        end
      end

      @session.init_ui(user_input, user_output) if @session
      @sysinfo = nil
    end

    # Meterpreter sometimes needs a little bit of extra time to
    # actually be responsive for post modules. Default tries
    # and retries for 5 seconds.
    def check_for_session_readiness(tries=6)
      session_ready_count = 0
      session_ready = false
      until session.sys or session_ready_count > tries
        session_ready_count += 1
        back_off_period = (session_ready_count**2)/10.0
        select(nil,nil,nil,back_off_period)
      end
      session_ready = !!session.sys
      unless session_ready
        raise "The stdapi extension has not been loaded yet." unless session.tlv_enc_key.nil?
        raise "Could not get a hold of the session."
      end
      return session_ready
    end

    #
    # Default cleanup handler does nothing
    #
    def cleanup
      super if defined?(super)
    end

    #
    # Return the associated session or nil if there isn't one
    #
    # @return [Msf::Session]
    # @return [nil] if the id provided in the datastore does not
    #   correspond to a session
    def session
      # Try the cached one
      return @session if @session && !session_changed?

      if datastore['SESSION']
        @session = framework.sessions.get(datastore['SESSION'].to_i)
      else
        @session = nil
      end

      @session
    end

    def session_display_info
      "Session: #{session.sid} (#{session.session_host})"
    end

    alias :client :session

    #
    # Cached sysinfo, returns nil for non-meterpreter sessions
    #
    # @return [Hash,nil]
    def sysinfo
      begin
        @sysinfo ||= session.sys.config.sysinfo
      rescue NoMethodError
        @sysinfo = nil
      end
      @sysinfo
    end

    #
    # Can be overridden by individual modules to add new commands
    #
    def post_commands
      {}
    end

    # Whether this module's {Msf::Exploit::Stance} is {Msf::Exploit::Stance::Passive passive}
    def passive?
      passive
    end

    #
    # Return a (possibly empty) list of all compatible sessions
    #
    # @return [Array]
    def compatible_sessions
      sessions = []
      framework.sessions.each do |sid, s|
        sessions << sid if session_compatible?(s)
      end
      sessions
    end

    #
    # Return false if the given session is not compatible with this module
    #
    # Checks the session's type against this module's
    # <tt>module_info["SessionTypes"]</tt> as well as examining platform
    # and arch compatibility.
    #
    # +sess_or_sid+ can be a Session object, Integer, or
    # String. In the latter cases it should be a key in
    # +framework.sessions+.
    #
    # @note Because it errs on the side of compatibility, a true return
    #   value from this method does not guarantee the module will work
    #   with the session. For example, ARCH_CMD modules can work on a
    #   variety of platforms and archs and thus return true in this check.
    #
    # @param sess_or_sid [Msf::Session,Integer,String]
    #   A session or session ID to compare against this module for
    #   compatibility.
    #
    def session_compatible?(sess_or_sid)
      session_incompatibility_reasons(sess_or_sid).empty?
    end

    #
    # Return the reasons why a session is incompatible.
    #
    # @return Array<String>
    def session_incompatibility_reasons(sess_or_sid)
      # Normalize the argument to an actual Session
      case sess_or_sid
      when ::Integer, ::String
        s = framework.sessions[sess_or_sid.to_i]
      when ::Msf::Session
        s = sess_or_sid
      when nil
        # No session provided
        return []
      end

      issues = []

      # Can't do anything without a session
      unless s
        issues << ['invalid session']
        return issues
      end

      # Can't be compatible if it's the wrong type
      if session_types && !session_types.include?(s.type)
        issues << "incompatible session type: #{s.type}. This module works with: #{session_types.join(', ')}."
      end

      # Check to make sure architectures match
      mod_arch = module_info['Arch']

      if mod_arch
        if s.arch.blank?
          issues << 'Unknown session arch'
        else
          mod_arch = Array.wrap(mod_arch)
          # Assume ARCH_CMD modules can work on supported SessionTypes since both shell and meterpreter types can execute commands
          issues << "incompatible session architecture: #{s.arch}" unless mod_arch.include?(s.arch) || mod_arch.include?(ARCH_CMD)
        end
      end

      # Arch is okay, now check the platform.

      if platform && platform.is_a?(Msf::Module::PlatformList) && !platform.empty?
        if s.platform.blank?
          issues << "Unknown session platform. This module works with: #{platform.names.join(', ')}."
        elsif !platform.supports?(Msf::Module::PlatformList.transform(s.platform))
          issues << "incompatible session platform: #{s.platform}. This module works with: #{platform ? platform.names.join(', ') : platform.inspect}."
        end
      end

      # Check all specified meterpreter commands are provided by the remote session
      if s.type == 'meterpreter'
        issues += meterpreter_session_incompatibility_reasons(s)
      end

      issues
    end

    #
    # True when this module is passive, false when active
    #
    # @return [Boolean]
    # @see passive?
    attr_reader :passive

    #
    # A list of compatible session types
    #
    # @return [Array]
    attr_reader :session_types

    protected

    attr_writer :passive, :session_types

    def session_changed?
      @ds_session ||= datastore['SESSION']

      if (@ds_session != datastore['SESSION'])
        @ds_session = nil
        return true
      else
        return false
      end
    end

    #
    # Return the reasons why a meterpreter session is incompatible. Checks all specified meterpreter commands
    # are provided by the remote session
    #
    # @return Array<String>
    def meterpreter_session_incompatibility_reasons(session)
      cmd_name_wildcards = module_info.dig('Compat', 'Meterpreter', 'Commands') || []
      cmd_names = Rex::Post::Meterpreter::CommandMapper.get_command_names.select do |cmd_name|
        cmd_name_wildcards.any? { |cmd_name_wildcard| ::File.fnmatch(cmd_name_wildcard, cmd_name) }
      end

      unmatched_wildcards = cmd_name_wildcards.select { |cmd_name_wildcard| cmd_names.none? { |cmd_name| ::File.fnmatch(cmd_name_wildcard, cmd_name) } }
      unless unmatched_wildcards.empty?
        # This implies that there was a typo in one of the wildcards because it didn't match anything. This is a developer mistake.
        wlog("The #{fullname} module specified the following Meterpreter command wildcards that did not match anything: #{ unmatched_wildcards.join(', ') }")
      end

      cmd_ids = cmd_names.map { |name| Rex::Post::Meterpreter::CommandMapper.get_command_id(name) }

      # XXX: Remove this condition once the payloads gem has had another major version bump from 2.x to 3.x and
      # rapid7/metasploit-payloads#451 has been landed to correct the `enumextcmd` behavior on Windows. Until then, skip
      # proactive validation of Windows core commands. This is not the only instance of this workaround.
      if session.base_platform == 'windows'
        cmd_ids = cmd_ids.select do |cmd_id|
          !cmd_id.between?(
            Rex::Post::Meterpreter::ClientCore.extension_id,
            Rex::Post::Meterpreter::ClientCore.extension_id + Rex::Post::Meterpreter::COMMAND_ID_RANGE - 1
          )
        end
      end

      # Windows does not support chmod, but will be defined by default in the file mixin
      if session.base_platform == 'windows'
        cmd_ids -= [Rex::Post::Meterpreter::Extensions::Stdapi::COMMAND_ID_STDAPI_FS_CHMOD]
      end

      missing_cmd_ids = (cmd_ids - session.commands)
      unless missing_cmd_ids.empty?
        # If there are missing commands, try to load the necessary extension.

        # If core_loadlib isn't supported, then extensions can't be loaded
        return ['missing Meterpreter features: core can not be extended'] unless session.commands.include?(Rex::Post::Meterpreter::COMMAND_ID_CORE_LOADLIB)

        # Since core is already loaded, if the missing command is a core command then it's truly missing
        missing_core_cmd_ids = missing_cmd_ids.select do |cmd_id|
          cmd_id.between?(
            Rex::Post::Meterpreter::ClientCore.extension_id,
            Rex::Post::Meterpreter::ClientCore.extension_id + Rex::Post::Meterpreter::COMMAND_ID_RANGE - 1
          )
        end
        if missing_core_cmd_ids.any?
          return ["missing Meterpreter features: #{command_names_for(missing_core_cmd_ids)}"]
        end

        missing_extensions = missing_cmd_ids.map { |cmd_id| Rex::Post::Meterpreter::ExtensionMapper.get_extension_name(cmd_id) }.uniq
        missing_extensions.each do |ext_name|
          # If the extension is already loaded, the command is truly missing
          return ["missing Meterpreter features: #{command_names_for(missing_cmd_ids)}"] if session.ext.aliases.include?(ext_name)

          begin
            session.core.use(ext_name)
          rescue RuntimeError
            return ["unloadable Meterpreter extension: #{ext_name}"]
          end
        end
      end
      missing_cmd_ids -= session.commands
      return ["missing Meterpreter features: #{command_names_for(missing_cmd_ids)}"] unless missing_cmd_ids.empty?

      []
    end

    def command_names_for(command_ids)
      command_ids.map { |id| Rex::Post::Meterpreter::CommandMapper.get_command_name(id) }.join(', ')
    end

  end
end