SQL sessions have correct history manager support

Author: sjanusz-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.
@@ -235,20 +243,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_interactive_history
+    self.new.postgresql_session_interactive_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 MSSQL interactive query history file
+  #
+  # @return [String] path to the interactive query history file.
+  def self.mssql_session_interactive_history
+    self.new.mssql_session_interactive_history
+  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_interactive_history
+    self.new.mysql_session_interactive_history
+  end
+
   def self.pry_history
     self.new.pry_history
   end
@@ -336,6 +365,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.
@@ -363,15 +399,27 @@ def ldap_session_history
   end
 
   def postgresql_session_history
-    config_directory + FileSep + "postgresql_session_history"
+    histories_directory + FileSep + "postgresql_session_history"
+  end
+
+  def postgresql_session_interactive_history
+    histories_directory + FileSep + "postgresql_session_interactive_history"
   end
 
   def mysql_session_history
-    config_directory + FileSep + "mysql_session_history"
+    histories_directory + FileSep + "mysql_session_history"
+  end
+
+  def mysql_session_interactive_history
+    histories_directory + FileSep + "mysql_session_interactive_history"
   end
 
   def mssql_session_history
-    config_directory + FileSep + "mssql_session_history"
+    histories_directory + FileSep + "mssql_session_history"
+  end
+
+  def mssql_session_interactive_history
+    histories_directory + FileSep + "mssql_session_interactive_history"
   end
 
   def pry_history
@@ -495,6 +543,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/sql/ui/console/interactive_sql_client.rb

@@ -67,8 +67,20 @@ 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.send("#{name}_session_interactive_history")
+
+    # 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 +170,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,21 @@ def _debug=(value)
     @debug = value
   end
 
+  # Return a queue of threads that have not yet finished saving the history file to disk.
+  # @return [Queue] a queue of threads that have not yet finished saving the history file to disk.
+  def _remaining_work
+    @remaining_work
+  end
+
   private
 
   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 +99,78 @@ 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
-      end
+    ::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
+
+    begin
+      File.open(history_file, 'r') do |f|
+        context[:input_library] == :reline ? clear_reline : clear_readline
+        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
+    rescue Errno::EACCES, Errno::ENOENT => e
+      $stderr.puts("Failed to open history file: #{history_file} with error: #{e}") if debug?
     end
   end
 
-  def store_history_file(history_file)
-    return unless readline_available?
+  def store_history_file(context)
+    history_file = context[:history_file]
+    history = context[:input_library] == :reline ? ::Reline::HISTORY : ::Readline::HISTORY
+
+    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)
@@ -144,7 +183,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