Land rapid7#18933, update SQL sessions to correctly manage history

Author: adfoster-r7

lib/msf/base/config.rb

@@ -228,20 +228,41 @@ def self.postgresql_session_history
     self.new.postgresql_session_history
   end
 
+  # Returns the full path to the PostgreSQL interactive query history file
+  #
+  # @return [String] path to the interactive query history file.
+  def self.postgresql_session_history_interactive
+    self.new.postgresql_session_history_interactive
+  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 MSSQL interactive query history file
+  #
+  # @return [String] path to the interactive query history file.
+  def self.mssql_session_history_interactive
+    self.new.mssql_session_history_interactive
+  end
+
   # Returns the full path to the MySQL session history file.
   #
   # @return [String] path to the history file.
   def self.mysql_session_history
     self.new.mysql_session_history
   end
 
+  # Returns the full path to the MySQL interactive query history file
+  #
+  # @return [String] path to the interactive query history file.
+  def self.mysql_session_history_interactive
+    self.new.mysql_session_history_interactive
+  end
+
   def self.pry_history
     self.new.pry_history
   end
@@ -355,14 +376,26 @@ def postgresql_session_history
     config_directory + FileSep + "postgresql_session_history"
   end
 
+  def postgresql_session_history_interactive
+    postgresql_session_history + "_interactive"
+  end
+
   def mysql_session_history
     config_directory + FileSep + "mysql_session_history"
   end
 
+  def mysql_session_history_interactive
+    mysql_session_history + "_interactive"
+  end
+
   def mssql_session_history
     config_directory + FileSep + "mssql_session_history"
   end
 
+  def mssql_session_history_interactive
+    mssql_session_history + "_interactive"
+  end
+
   def pry_history
     config_directory + FileSep + "pry_history"
   end

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

@@ -67,8 +67,19 @@ 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 = {}
+
+    # Multiline (Reline) and fallback (Readline) have separate history contexts as they are two different libraries.
+    framework.history_manager.with_context(history_file: Msf::Config.send("#{name}_session_history_interactive"), name: name, input_library: :reline) do
+      query = _multiline
+    end
+
+    if query[:status] == :fail
+      framework.history_manager.with_context(history_file: Msf::Config.send("#{name}_session_history_interactive"), name: name, input_library: :readline) do
+        query = _fallback
+      end
+    end
 
     query
   end
@@ -163,6 +174,16 @@ def _fallback
 
   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
@@ -65,9 +67,9 @@ def debug?
     @debug
   end
 
-  def push_context(history_file: nil, name: nil)
+  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,47 +93,69 @@ 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?
+
+    ::Reline::HISTORY.length.times { ::Reline::HISTORY.pop }
+  end
+
+  def load_history_file(context)
+    history_file = context[:history_file]
+    history = context[:input_library] == :reline ? ::Reline::HISTORY : ::Readline::HISTORY
 
-    clear_readline
     if File.exist?(history_file)
-      File.readlines(history_file).each do |e|
-        ::Readline::HISTORY << e.chomp
+      File.open(history_file, 'r') do |f|
+        f.each do |line|
+          chomped_line = line.chomp
+          if context[:input_library] == :reline && history.last&.end_with?("\\")
+            history.last.delete_suffix!("\\")
+            history.last << "\n" << chomped_line
+          else
+            history << chomped_line
+          end
+        end
       end
     end
   end
 
-  def store_history_file(history_file)
-    return unless readline_available?
-    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?
-    end
+  def store_history_file(context)
+    history_file = context[:history_file]
+    history = context[:input_library] == :reline ? ::Reline::HISTORY : ::Readline::HISTORY
+
+    history_to_save = history.map { |line| line.scrub.split("\n").join("\\\n") }
 
-    write_history_file(history_file, cmds)
+    write_history_file(history_file, history_to_save)
   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)
@@ -144,7 +168,7 @@ def write_history_file(history_file, cmds)
           cmds = event[:cmds]
 
           File.open(history_file, 'wb+') do |f|
-            f.puts(cmds.reverse)
+            f.puts(cmds)
           end
 
         rescue => e

spec/lib/rex/ui/text/shell/history_manager_spec.rb

@@ -25,7 +25,7 @@
         (expect do |block|
           subject.with_context(name: 'a') do
             expected_contexts = [
-              { history_file: nil, name: 'a' },
+              { history_file: nil, input_library: :readline, name: 'a' },
             ]
             expect(subject._contexts).to eq(expected_contexts)
             block.to_proc.call
@@ -36,15 +36,15 @@
 
     context 'when there is an existing stack' do
       before(:each) do
-        subject.send(:push_context, history_file: nil, name: 'a')
+        subject.send(:push_context, history_file: nil, input_library: :readline, name: 'a')
       end
 
       it 'continues to have the previous existing stack' do
         subject.with_context {
           # noop
         }
         expected_contexts = [
-          { history_file: nil, name: 'a' },
+          { history_file: nil, input_library: :readline, name: 'a' },
         ]
         expect(subject._contexts).to eq(expected_contexts)
       end
@@ -53,8 +53,8 @@
         (expect do |block|
           subject.with_context(name: 'b') do
             expected_contexts = [
-              { history_file: nil, name: 'a' },
-              { history_file: nil, name: 'b' },
+              { history_file: nil, input_library: :readline, name: 'a' },
+              { history_file: nil, input_library: :readline, name: 'b' },
             ]
             expect(subject._contexts).to eq(expected_contexts)
             block.to_proc.call
@@ -69,7 +69,7 @@
           }
         end.to raise_exception ArgumentError, 'Mock error'
         expected_contexts = [
-          { history_file: nil, name: 'a' },
+          { history_file: nil, input_library: :readline, name: 'a' },
         ]
         expect(subject._contexts).to eq(expected_contexts)
       end
@@ -79,9 +79,9 @@
   describe '#push_context' do
     context 'when the stack is empty' do
       it 'stores the history contexts' do
-        subject.send(:push_context, history_file: nil, name: 'a')
+        subject.send(:push_context, history_file: nil, input_library: :readline, name: 'a')
         expected_contexts = [
-          { history_file: nil, name: 'a' }
+          { history_file: nil, input_library: :readline, name: 'a' }
         ]
         expect(subject._contexts).to eq(expected_contexts)
       end
@@ -90,12 +90,12 @@
     context 'when multiple values are pushed' do
       it 'stores the history contexts' do
         subject.send(:push_context, history_file: nil, name: 'a')
-        subject.send(:push_context, history_file: nil, name: 'b')
-        subject.send(:push_context, history_file: nil, name: 'c')
+        subject.send(:push_context, history_file: nil, input_library: :readline, name: 'b')
+        subject.send(:push_context, history_file: nil, input_library: :reline, name: 'c')
         expected_contexts = [
-          { history_file: nil, name: 'a' },
-          { history_file: nil, name: 'b' },
-          { history_file: nil, name: 'c' },
+          { history_file: nil, input_library: :readline, name: 'a' },
+          { history_file: nil, input_library: :readline, name: 'b' },
+          { history_file: nil, input_library: :reline, name: 'c' },
         ]
         expect(subject._contexts).to eq(expected_contexts)
       end
@@ -113,12 +113,12 @@
     end
 
     context 'when the stack is not empty' do
-      it 'continues to have an empty stack' do
+      it 'continues to have a non-empty stack' do
         subject.send(:push_context, history_file: nil, name: 'a')
         subject.send(:push_context, history_file: nil, name: 'b')
         subject.send(:pop_context)
         expected_contexts = [
-          { history_file: nil, name: 'a' },
+          { history_file: nil, input_library: :readline, name: 'a' },
         ]
         expect(subject._contexts).to eq(expected_contexts)
       end