Land rapid7#19197, SQL sessions have correct history manager support

Author: adfoster-r7

lib/msf/base/config.rb

@@ -74,7 +74,8 @@ def self.get_config_root
       'PluginDirectory'     => "plugins",
       'DataDirectory'       => "data",
       'LootDirectory'       => "loot",
-      'LocalDirectory'      => "local"
+      'LocalDirectory'      => "local",
+      'HistoriesDirectory'  => "histories"
     }
 
   ##
@@ -97,6 +98,13 @@ def self.config_directory
     self.new.config_directory
   end
 
+  # Returns the histories directory default.
+  #
+  # @return [String] the SQL session histories directory.
+  def self.histories_directory
+    self.new.histories_directory
+  end
+
   # Return the directory that logo files should be loaded from.
   #
   # @return [String] path to the logos directory.
@@ -228,25 +236,11 @@ def self.ldap_session_history
     self.new.ldap_session_history
   end
 
-  # Returns the full path to the PostgreSQL session history file.
-  #
-  # @return [String] path to the history file.
-  def self.postgresql_session_history
-    self.new.postgresql_session_history
-  end
-
-  # Returns the full path to the MSSQL session history file.
-  #
-  # @return [String] path to the history file.
-  def self.mssql_session_history
-    self.new.mssql_session_history
-  end
-
-  # Returns the full path to the MySQL session history file.
+  # Returns the full path to the MySQL interactive query history file
   #
-  # @return [String] path to the history file.
-  def self.mysql_session_history
-    self.new.mysql_session_history
+  # @return [String] path to the interactive query history file.
+  def self.history_file_for_session_type(opts)
+    self.new.history_file_for_session_type(opts)
   end
 
   def self.pry_history
@@ -336,6 +330,13 @@ def config_directory
     self['ConfigDirectory']
   end
 
+  # Returns the histories directory default.
+  #
+  # @return [String] the SQL session histories directory.
+  def histories_directory
+    config_directory + FileSep + self['HistoriesDirectory']
+  end
+
   # Returns the full path to the configuration file.
   #
   # @return [String] path to the configuration file.
@@ -362,16 +363,24 @@ def ldap_session_history
     config_directory + FileSep + "ldap_session_history"
   end
 
-  def postgresql_session_history
-    config_directory + FileSep + "postgresql_session_history"
+  def history_options_valid?(opts)
+    return false if (opts[:session_type].nil? || opts[:interactive].nil?)
+
+    true
   end
 
-  def mysql_session_history
-    config_directory + FileSep + "mysql_session_history"
+  def interactive_to_string_map(interactive)
+    # Check for true explicitly rather than just a value that is truthy.
+    interactive == true ? '_interactive' : ''
   end
 
-  def mssql_session_history
-    config_directory + FileSep + "mssql_session_history"
+  def history_file_for_session_type(opts)
+    return nil unless history_options_valid?(opts)
+
+    session_type_name = opts[:session_type]
+    interactive = interactive_to_string_map(opts[:interactive])
+
+    histories_directory + FileSep + "#{session_type_name}_session#{interactive}_history"
   end
 
   def pry_history
@@ -495,6 +504,7 @@ def init
     FileUtils.mkdir_p(user_module_directory)
     FileUtils.mkdir_p(user_plugin_directory)
     FileUtils.mkdir_p(user_data_directory)
+    FileUtils.mkdir_p(histories_directory)
   end
 
   # Loads configuration from the supplied file path, or the default one if

lib/rex/post/mssql/ui/console.rb

@@ -30,8 +30,8 @@ def initialize(session, opts={})
             self.client = session.client
             envchange = ::Rex::Proto::MSSQL::ClientMixin::ENVCHANGE
             prompt = "%undMSSQL @ #{client.peerinfo} (#{client.initial_info_for_envchange(envchange: envchange::DATABASE)[:new]})%clr"
-            history_manager = Msf::Config.mssql_session_history
-            super(prompt, '>', history_manager, nil, :mssql)
+            history_file = Msf::Config.history_file_for_session_type(session_type: session.type, interactive: false)
+            super(prompt, '>', history_file, nil, :mssql)
 
             # Queued commands array
             self.commands = []

lib/rex/post/mysql/ui/console.rb

@@ -25,8 +25,8 @@ def initialize(session)
             self.session = session
             self.client = session.client
             prompt = "%undMySQL @ #{client.peerinfo} (#{current_database})%clr"
-            history_manager = Msf::Config.mysql_session_history
-            super(prompt, '>', history_manager, nil, :mysql)
+            history_file = Msf::Config.history_file_for_session_type(session_type: session.type, interactive: false)
+            super(prompt, '>', history_file, nil, :mysql)
 
             # Queued commands array
             self.commands = []

lib/rex/post/postgresql/ui/console.rb

@@ -29,8 +29,8 @@ def initialize(session)
             self.session = session
             self.client = session.client
             prompt = "%undPostgreSQL @ #{client.peerinfo} (#{current_database})%clr"
-            history_manager = Msf::Config.postgresql_session_history
-            super(prompt, '>', history_manager, nil, :postgresql)
+            history_file = Msf::Config.history_file_for_session_type(session_type: session.type, interactive: false)
+            super(prompt, '>', history_file, nil, :postgresql)
 
             # Queued commands array
             self.commands = []

lib/rex/post/sql/ui/console/interactive_sql_client.rb

@@ -67,8 +67,21 @@ def _winch
 
   # Try getting multi-line input support provided by Reline, fall back to Readline.
   def _multiline_with_fallback
-    query = _multiline
-    query = _fallback if query[:status] == :fail
+    name = session.type
+    query = {}
+    history_file = Msf::Config.history_file_for_session_type(session_type: name, interactive: true)
+    return { status: :fail, errors: ["Unable to get history file for session type: #{name}"] } if history_file.nil?
+
+    # Multiline (Reline) and fallback (Readline) have separate history contexts as they are two different libraries.
+    framework.history_manager.with_context(history_file: history_file , name: name, input_library: :reline) do
+      query = _multiline
+    end
+
+    if query[:status] == :fail
+      framework.history_manager.with_context(history_file: history_file, name: name, input_library: :readline) do
+        query = _fallback
+      end
+    end
 
     query
   end
@@ -158,11 +171,21 @@ def _fallback
       break if line.end_with? ';'
     end
 
-    { status: :success, result: line_buffer.join }
+    { status: :success, result: line_buffer.join(' ') }
   end
 
   attr_accessor :on_log_proc, :client_dispatcher
 
+  private
+
+  def framework
+    client_dispatcher.shell.framework
+  end
+
+  def session
+    client_dispatcher.shell.session
+  end
+
 end
 end
 end

lib/rex/ui/text/shell/history_manager.rb

@@ -24,10 +24,12 @@ def initialize
   #
   # @param [String,nil] history_file The file to load and persist commands to
   # @param [String] name Human readable history context name
+  # @param [Symbol] input_library The input library to provide context for. :reline, :readline
   # @param [Proc] block
   # @return [nil]
-  def with_context(history_file: nil, name: nil, &block)
-    push_context(history_file: history_file, name: name)
+  def with_context(history_file: nil, name: nil, input_library: nil, &block)
+    # Default to Readline for backwards compatibility.
+    push_context(history_file: history_file, name: name, input_library: input_library || :readline)
 
     begin
       block.call
@@ -59,15 +61,45 @@ def _debug=(value)
     @debug = value
   end
 
+  def _close
+    event = { type: :close }
+    @write_queue << event
+    @remaining_work << event
+  end
+
   private
 
   def debug?
     @debug
   end
 
-  def push_context(history_file: nil, name: nil)
+  # A wrapper around mapping the input library to its history; this way we can mock the return value of this method.
+  def map_library_to_history(input_library)
+    case input_library
+    when :readline
+      ::Readline::HISTORY
+    when :reline
+      ::Reline::HISTORY
+    else
+      $stderr.puts("Unknown input library: #{input_library}") if debug?
+      []
+    end
+  end
+
+  def clear_library(input_library)
+    case input_library
+    when :readline
+      clear_readline
+    when :reline
+      clear_reline
+    else
+      $stderr.puts("Unknown input library: #{input_library}") if debug?
+    end
+  end
+
+  def push_context(history_file: nil, name: nil, input_library: nil)
     $stderr.puts("Push context before\n#{JSON.pretty_generate(_contexts)}") if debug?
-    new_context = { history_file: history_file, name: name }
+    new_context = { history_file: history_file, name: name, input_library: input_library || :readline }
 
     switch_context(new_context, @contexts.last)
     @contexts.push(new_context)
@@ -91,60 +123,93 @@ def readline_available?
     defined?(::Readline)
   end
 
+  def reline_available?
+    begin
+      require 'reline'
+      defined?(::Reline)
+    rescue ::LoadError => _e
+      false
+    end
+  end
+
   def clear_readline
     return unless readline_available?
 
     ::Readline::HISTORY.length.times { ::Readline::HISTORY.pop }
   end
 
-  def load_history_file(history_file)
-    return unless readline_available?
+  def clear_reline
+    return unless reline_available?
 
-    clear_readline
-    if File.exist?(history_file)
-      File.readlines(history_file).each do |e|
-        ::Readline::HISTORY << e.chomp
+    ::Reline::HISTORY.length.times { ::Reline::HISTORY.pop }
+  end
+
+  def load_history_file(context)
+    history_file = context[:history_file]
+    history = map_library_to_history(context[:input_library])
+
+    begin
+      File.open(history_file, 'rb') do |f|
+        clear_library(context[:input_library])
+        f.each_line(chomp: true) do |line|
+          if context[:input_library] == :reline && history.last&.end_with?("\\")
+            history.last.delete_suffix!("\\")
+            history.last << "\n" << line
+          else
+            history << line
+          end
+        end
       end
+    rescue Errno::EACCES, Errno::ENOENT => e
+      elog "Failed to open history file: #{history_file} with error: #{e}"
     end
   end
 
-  def store_history_file(history_file)
-    return unless readline_available?
+  def store_history_file(context)
+    history_file = context[:history_file]
+    history = map_library_to_history(context[:input_library])
+
+    history_diff = history.length < MAX_HISTORY ? history.length : MAX_HISTORY
+
     cmds = []
-    history_diff = ::Readline::HISTORY.length < MAX_HISTORY ? ::Readline::HISTORY.length : MAX_HISTORY
     history_diff.times do
-      entry = ::Readline::HISTORY.pop
-      cmds.push(entry) unless entry.nil?
+      entry = history.pop
+      cmds << entry.scrub.split("\n").join("\\\n")
     end
 
-    write_history_file(history_file, cmds)
+    write_history_file(history_file, cmds.reverse)
   end
 
   def switch_context(new_context, old_context=nil)
     if old_context && old_context[:history_file]
-      store_history_file(old_context[:history_file])
+      store_history_file(old_context)
     end
 
     if new_context && new_context[:history_file]
-      load_history_file(new_context[:history_file])
+      load_history_file(new_context)
     else
       clear_readline
+      clear_reline
     end
-  rescue SignalException => e
+  rescue SignalException => _e
     clear_readline
+    clear_reline
   end
 
   def write_history_file(history_file, cmds)
     write_queue_ref = @write_queue
     remaining_work_ref = @remaining_work
+
     @write_thread ||= Rex::ThreadFactory.spawn("HistoryManagerWriter", false) do
       while (event = write_queue_ref.pop)
         begin
+          break if event[:type] == :close
+
           history_file = event[:history_file]
           cmds = event[:cmds]
 
           File.open(history_file, 'wb+') do |f|
-            f.puts(cmds.reverse)
+            f.puts(cmds)
           end
 
         rescue => e