Support CTRL + Z for interactive SQL REPL

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

@@ -18,6 +18,10 @@ class Console
           require 'rex/post/postgresql/ui/console/command_dispatcher/client'
           require 'rex/post/postgresql/ui/console/command_dispatcher/modules'
 
+          # Interactive channel, required for the REPL shell interaction and correct CTRL + Z handling.
+          # Zeitwerk ignored `rex/post` files so we need to `require` this file here.
+          require 'rex/post/postgresql/ui/console/interactive_channel'
+
           #
           # Initialize the PostgreSQL console.
           #
@@ -109,6 +113,23 @@ def log_error(msg)
             dlog("Call stack:\n#{$@.join("\n")}", 'postgresql')
           end
 
+          #
+          # Interacts with the supplied client.
+          # TODO: Is there a better name than interact_with_client, as it's not quite what we want
+          #
+          def interact_with_client(client_dispatcher: nil)
+            return unless client_dispatcher
+
+            client.extend(InteractiveChannel) unless (client.kind_of?(InteractiveChannel) == true)
+            client.on_command_proc = self.on_command_proc if self.on_command_proc
+            client.on_print_proc   = self.on_print_proc if self.on_print_proc
+            client.on_log_proc = method(:log_output) if self.respond_to?(:log_output, true)
+            client.client_dispatcher = client_dispatcher
+
+            client.interact(input, output)
+            client.reset_ui
+          end
+
           # @return [Msf::Sessions::PostgreSQL]
           attr_reader :session
 

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

@@ -66,42 +66,9 @@ def cmd_shell(*args)
               return
             end
 
-            stop_words = %w[stop s exit e end quit q].freeze
-
-            # Allow the user to query the DB in a loop.
-            finished = false
-            until finished
-              begin
-                # This needs to be here, otherwise the `ensure` block would reset it to the previous
-                # value after a single query, meaning future queries would have the default prompt_block.
-                prompt_proc_before = ::Reline.prompt_proc
-                ::Reline.prompt_proc = proc { |line_buffer| line_buffer.each_with_index.map { |_line, i| i > 0 ? 'SQL *> ' : 'SQL >> ' } }
-
-                # This will loop until it receives `true`.
-                raw_query = ::Reline.readmultiline('SQL >> ', use_history = true) do |multiline_input|
-                  # In the case only a stop word was input, exit out of the REPL shell
-                  finished = multiline_input.split.count == 1 && stop_words.include?(multiline_input.split.last)
-                  # Accept the input until the current line does not end with '\', similar to a shell
-                  finished || multiline_input.split.empty? || !multiline_input.split.last&.end_with?('\\')
-                end
-              rescue ::Interrupt => _e
-                finished = true
-              ensure
-                ::Reline.prompt_proc = prompt_proc_before
-              end
-
-              if finished
-                print_status 'Exiting Shell mode.'
-                return
-              end
-
-              formatted_query = process_query(query: raw_query)
-
-              unless formatted_query.empty?
-                print_status "Running SQL Command: '#{formatted_query}'"
-                cmd_query(formatted_query)
-              end
-            end
+            console = shell
+            # Pass in self so that we can call cmd_query in subsequent calls
+            console.interact_with_client(client_dispatcher: self)
           end
 
           def cmd_query_help
@@ -159,7 +126,7 @@ def cmd_query(*args)
           def process_query(query: '')
             return '' if query.empty?
 
-            query.lines.each.map { |line| line.chomp("\\\n").strip }.reject(&:empty?).compact.join(' ')
+            query.lines.each.map { |line| line.chomp.chomp('\\').strip }.reject(&:empty?).compact.join(' ')
           end
         end
       end

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

@@ -0,0 +1,160 @@
+# -*- coding: binary -*-
+module Rex
+module Post
+module PostgreSQL
+module Ui
+
+###
+#
+# Mixin that is meant to extend the base channel class from meterpreter in a
+# manner that adds interactive capabilities.
+#
+###
+module Console::InteractiveChannel
+
+  include Rex::Ui::Interactive
+
+  #
+  # Interacts with self.
+  #
+  def _interact
+    begin
+      while self.interacting
+        raw_sql_string = _get_user_input
+        # We need to check that the user is still interacting, i.e. if ctrl+z is triggered when requesting user input
+        break unless (self.interacting && raw_sql_string)
+
+        formatted_query = client_dispatcher.process_query(query: raw_sql_string)
+        print_status "Executing query: #{formatted_query}"
+        client_dispatcher.cmd_query(formatted_query)
+      end
+    rescue ::StandardError => e
+      elog('Error when interact with SQL REPL', e)
+      raise e
+    end
+  end
+
+  #
+  # Called when an interrupt is sent.
+  #
+  def _interrupt
+    prompt_yesno('Terminate interactive REPL interpreter?')
+  end
+
+  #
+  # Suspends interaction with the interactive REPL interpreter
+  #
+  def _suspend
+    if (prompt_yesno('Background interactive REPL interpreter?') == true)
+      self.interacting = false
+    end
+  end
+
+  #
+  # We don't need to do any clean-up when finishing the interaction with the REPL
+  #
+  def _interact_complete
+    # noop
+  end
+
+  # To have this be consistent between reline and the fallback, this will return a single string only.
+  # Try getting multi-line input support provided by Reline, fall back to Readline.
+  def _multiline_with_fallback
+    begin
+      query = _multiline
+      query = _fallback if query[:status] == :fail
+    rescue ::StandardError => e
+      elog('Failed to get SQL input from user', e)
+      raise e
+    end
+
+    query[:result]
+  end
+
+  def _multiline
+    begin
+      require 'reline' unless defined?(::Reline)
+      user_input.extend(::Reline) unless user_input.kind_of?(::Reline)
+    rescue ::LoadError => e
+      elog('Failed to load Reline', e)
+      return { status: :fail, errors: [e] }
+    end
+
+    stop_words = %w[stop s exit e end quit q].freeze
+
+    finished = false
+    begin
+      prompt_proc_before = ::Reline.prompt_proc
+      ::Reline.prompt_proc = proc { |line_buffer| line_buffer.each_with_index.map { |_line, i| i > 0 ? 'SQL *> ' : 'SQL >> ' } }
+
+      # We want to do this in a loop
+      raw_query = user_input.send(:readmultiline, 'SQL >> ', use_history = true) do |multiline_input|
+        # The user pressed ctrl + c or ctrl + z and wants to background our SQL prompt
+        return { status: :exit, result: nil } unless self.interacting
+
+        # In the case only a stop word was input, exit out of the REPL shell
+        finished = (multiline_input.split.count == 1 && stop_words.include?(multiline_input.split.last))
+
+        finished || multiline_input.split.last&.end_with?(';')
+      end
+    rescue ::StandardError => e
+      elog('Failed to get multi-line SQL query from user', e)
+    ensure
+      ::Reline.prompt_proc = prompt_proc_before
+    end
+
+    if finished
+      print_status 'Exiting Shell mode.'
+      self.interacting = false
+      return { status: :exit, result: nil }
+    end
+
+    { status: :success, result: raw_query }
+  end
+
+  def _fallback
+    stop_words = %w[stop s exit e end quit q].freeze
+    line_buffer = []
+    while (line = ::Readline.readline(prompt = line_buffer.empty? ? 'SQL >> ' : 'SQL *> ', add_history = true))
+      return { status: :exit, result: nil } unless self.interacting
+
+      if stop_words.include? line.chomp.downcase
+        print_status 'Exiting shell mode.'
+        self.interacting = false
+        return { status: :exit, result: nil }
+      end
+
+      next if line.empty?
+
+      line_buffer.append line
+
+      break if line.end_with? ';'
+    end
+
+    { status: :success, result: line_buffer.join }
+  end
+
+  #
+  # Reads data from local input
+  #
+  def _get_user_input
+    data = _multiline_with_fallback
+    return unless data
+
+    self.on_command_proc.call(data.strip) if self.on_command_proc
+
+    data
+  end
+
+  def _winch
+    # noop ?
+  end
+
+  attr_accessor :on_log_proc, :client_dispatcher
+
+end
+
+end
+end
+end
+end