Add _close method to history manager, use generic method for Msf Config history file

Author: sjanusz-r7

lib/msf/base/config.rb

@@ -236,46 +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 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
+  def self.history_file_for_session_type(opts)
+    self.new.history_file_for_session_type(opts)
   end
 
   def self.pry_history
@@ -398,28 +363,24 @@ def ldap_session_history
     config_directory + FileSep + "ldap_session_history"
   end
 
-  def postgresql_session_history
-    histories_directory + FileSep + "postgresql_session_history"
-  end
+  def history_options_valid?(opts)
+    return false if (opts[:session_type].nil? || opts[:interactive].nil?)
 
-  def postgresql_session_interactive_history
-    histories_directory + FileSep + "postgresql_session_interactive_history"
+    true
   end
 
-  def mysql_session_history
-    histories_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 mysql_session_interactive_history
-    histories_directory + FileSep + "mysql_session_interactive_history"
-  end
+  def history_file_for_session_type(opts)
+    return nil unless history_options_valid?(opts)
 
-  def mssql_session_history
-    histories_directory + FileSep + "mssql_session_history"
-  end
+    session_type_name = opts[:session_type]
+    interactive = interactive_to_string_map(opts[:interactive])
 
-  def mssql_session_interactive_history
-    histories_directory + FileSep + "mssql_session_interactive_history"
+    histories_directory + FileSep + "#{session_type_name}_session#{interactive}_history"
   end
 
   def pry_history

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

@@ -69,7 +69,8 @@ def _winch
   def _multiline_with_fallback
     name = session.type
     query = {}
-    history_file = Msf::Config.send("#{name}_session_interactive_history")
+    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

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

@@ -61,10 +61,10 @@ 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
+  def _close
+    event = { type: :close }
+    @write_queue << event
+    @remaining_work << event
   end
 
   private
@@ -149,20 +149,19 @@ def load_history_file(context)
     history = map_library_to_history(context[:input_library])
 
     begin
-      File.open(history_file, 'r') do |f|
+      File.open(history_file, 'rb') do |f|
         clear_library(context[:input_library])
-        f.each do |line|
-          chomped_line = line.chomp
+        f.each_line(chomp: true) do |line|
           if context[:input_library] == :reline && history.last&.end_with?("\\")
             history.last.delete_suffix!("\\")
-            history.last << "\n" << chomped_line
+            history.last << "\n" << line
           else
-            history << chomped_line
+            history << line
           end
         end
-        end
+      end
     rescue Errno::EACCES, Errno::ENOENT => e
-      $stderr.puts("Failed to open history file: #{history_file} with error: #{e}") if debug?
+      elog "Failed to open history file: #{history_file} with error: #{e}"
     end
   end
 
@@ -202,8 +201,10 @@ def write_history_file(history_file, cmds)
     remaining_work_ref = @remaining_work
 
     @write_thread ||= Rex::ThreadFactory.spawn("HistoryManagerWriter", false) do
-      while (write_queue_ref.size > 0 && event = write_queue_ref.pop)
+      while (event = write_queue_ref.pop)
         begin
+          break if event[:type] == :close
+
           history_file = event[:history_file]
           cmds = event[:cmds]
 

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

@@ -6,6 +6,8 @@
 require 'tempfile'
 
 RSpec.describe Rex::Ui::Text::Shell::HistoryManager do
+  include_context 'wait_for_expect'
+
   subject { described_class.send(:new) }
   let(:readline_available) { false }
   let(:reline_available) { false }
@@ -139,6 +141,7 @@
     after(:each) do
       # https://ruby-doc.org/stdlib-2.5.3/libdoc/tempfile/rdoc/Tempfile.html#class-Tempfile-label-Explicit+close
       history_file.unlink
+      subject._close
     end
 
     [
@@ -148,7 +151,6 @@
     ].each do |test|
       context "when storing #{test[:history_size]} lines" do
         it "correctly stores #{test[:expected_size]} lines" do
-          allow(subject).to receive(:_remaining_work).and_call_original
           allow(subject).to receive(:store_history_file).and_call_original
           allow(subject).to receive(:map_library_to_history).and_return(history_mock)
 
@@ -161,9 +163,9 @@
 
           subject.send(:store_history_file, context)
 
-          sleep(0.1) until subject._remaining_work.empty?
-
-          expect(history_file.read.split("\n").count).to eq(test[:expected_size])
+          wait_for_expect do
+            expect(history_file.read.split("\n").count).to eq(test[:expected_size])
+          end
         end
       end
     end
@@ -177,12 +179,13 @@
 
     after(:each) do
       history_file.unlink
+      subject._close
     end
 
     context 'when history file is not accessible' do
       it 'the library history remains unchanged' do
         allow(subject).to receive(:map_library_to_history).and_return(history_mock)
-        history_file = ::File.join('does/not/exist/history')
+        history_file = ::File.join('does', 'not', 'exist', 'history')
         context = { input_library: :readline, history_file: history_file, name: 'history' }
 
         subject.send(:load_history_file, context)