Add history support to nasm and metasm shells

lib/rex/ui/text/shell.rb

@@ -130,11 +130,8 @@ def run(&block)
       # Pry is a development dependency, if not available suppressing history_load can be safely ignored.
     end
 
-    framework.history_manager.with_context(history_file: histfile, name: name) do
-      self.hist_last_saved = Readline::HISTORY.length
-
+    with_history_manager_context do
       begin
-
         while true
           # If the stop flag was set or we've hit EOF, break out
           break if self.stop_flag || self.stop_count > 1
@@ -168,17 +165,13 @@ def run(&block)
             run_single(line)
             self.stop_count = 0
           end
-
         end
         # Prevent accidental console quits
       rescue ::Interrupt
         output.print("Interrupt: use the 'exit' command to quit\n")
         retry
       end
     end
-  ensure
-    framework.history_manager.flush
-    self.hist_last_saved = Readline::HISTORY.length
   end
 
   #
@@ -298,10 +291,29 @@ def print(msg='')
   attr_accessor :on_command_proc
   attr_accessor :on_print_proc
   attr_accessor :framework
+  attr_accessor :history_manager
   attr_accessor :hist_last_saved # the number of history lines when last saved/loaded
 
   protected
 
+  # Executes the yielded block under the context of a new HistoryManager context. The shell's history will be flushed
+  # to disk when no longer interacting with the shell. If no history manager is available, the history will not be persisted.
+  def with_history_manager_context
+    history_manager = self.history_manager || framework&.history_manager
+    return yield unless history_manager
+
+    begin
+      history_manager.with_context(history_file: histfile, name: name) do
+        self.hist_last_saved = Readline::HISTORY.length
+
+        yield
+      end
+    ensure
+      history_manager.flush
+      self.hist_last_saved = Readline::HISTORY.length
+    end
+  end
+
   def supports_color?
     true
   end

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

@@ -14,8 +14,10 @@ class HistoryManager
   def initialize
     @contexts = []
     @debug = false
+    # Values dequeued before work is started
     @write_queue = ::Queue.new
-    @currently_processing = ::Queue.new
+    # Values dequeued after work is completed
+    @remaining_work = ::Queue.new
   end
 
   # Create a new history command context when executing the given block
@@ -38,7 +40,7 @@ def with_context(history_file: nil, name: nil, &block)
 
   # Flush the contents of the write queue to disk. Blocks synchronously.
   def flush
-    until @write_queue.empty? && @currently_processing.empty?
+    until @write_queue.empty? && @remaining_work.empty?
       sleep 0.1
     end
 
@@ -134,27 +136,28 @@ def switch_context(new_context, old_context=nil)
 
   def write_history_file(history_file, cmds)
     write_queue_ref = @write_queue
-    currently_processing_ref = @currently_processing
+    remaining_work_ref = @remaining_work
     @write_thread ||= Rex::ThreadFactory.spawn("HistoryManagerWriter", false) do
       while (event = write_queue_ref.pop)
         begin
-          currently_processing_ref << event
-
           history_file = event[:history_file]
           cmds = event[:cmds]
 
           File.open(history_file, 'wb+') do |f|
             f.puts(cmds.reverse)
           end
+
         rescue => e
           elog(e)
         ensure
-          currently_processing_ref.pop
+          remaining_work_ref.pop
         end
       end
     end
 
-    write_queue_ref << { type: :write, history_file: history_file, cmds: cmds }
+    event = { type: :write, history_file: history_file, cmds: cmds }
+    @write_queue << event
+    @remaining_work << event
   end
 end
 

tools/exploit/metasm_shell.rb

@@ -160,8 +160,10 @@ def parse_gas_file(filename)
 
 # Start a pseudo shell and dispatch lines to be assembled and then
 # disassembled.
-shell = Rex::Ui::Text::PseudoShell.new("%bldmetasm%clr")
+history_file = File.join(Msf::Config.config_directory, 'metasm_history')
+shell = Rex::Ui::Text::PseudoShell.new("%bldmetasm%clr", '>', history_file)
 shell.init_ui(Rex::Ui::Text::Input::Stdio.new, Rex::Ui::Text::Output::Stdio.new)
+shell.history_manager = Rex::Ui::Text::Shell::HistoryManager.new
 
 puts [
   'type "exit" or "quit" to quit',

tools/exploit/nasm_shell.rb

@@ -39,8 +39,10 @@
 
 # Start a pseudo shell and dispatch lines to be assembled and then
 # disassembled.
-shell = Rex::Ui::Text::PseudoShell.new("%bldnasm%clr")
+history_file = File.join(Msf::Config.config_directory, 'nasm_history')
+shell = Rex::Ui::Text::PseudoShell.new("%bldnasm%clr", '>', history_file)
 shell.init_ui(Rex::Ui::Text::Input::Stdio.new, Rex::Ui::Text::Output::Stdio.new)
+shell.history_manager = Rex::Ui::Text::Shell::HistoryManager.new
 
 shell.run { |line|
   line.gsub!(/(\r|\n)/, '')