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_sql_client'
+
           #
           # Initialize the PostgreSQL console.
           #
@@ -109,6 +113,22 @@ def log_error(msg)
             dlog("Call stack:\n#{$@.join("\n")}", 'postgresql')
           end
 
+          #
+          # Interacts with the supplied client.
+          #
+          def interact_with_client(client_dispatcher: nil)
+            return unless client_dispatcher
+
+            client.extend(InteractiveSqlClient) unless (client.kind_of?(InteractiveSqlClient) == 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_sql_client.rb

@@ -0,0 +1,137 @@
+# -*- coding: binary -*-
+module Rex
+module Post
+module PostgreSQL
+module Ui
+
+###
+#
+# Mixin that is meant to extend a sql client class in a
+# manner that adds interactive capabilities.
+#
+###
+module Console::InteractiveSqlClient
+
+  include Rex::Ui::Interactive
+
+  #
+  # Interacts with self.
+  #
+  def _interact
+    while self.interacting
+      sql_input = _multiline_with_fallback
+      self.interacting = (sql_input[:status] != :exit)
+      # 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 && sql_input[:result])
+
+      self.on_command_proc.call(sql_input[:result].strip) if self.on_command_proc
+
+      formatted_query = client_dispatcher.process_query(query: sql_input[:result])
+      print_status "Executing query: #{formatted_query}"
+      client_dispatcher.cmd_query(formatted_query)
+    end
+  end
+
+  #
+  # Called when an interrupt is sent.
+  #
+  def _interrupt
+    prompt_yesno('Terminate interactive SQL prompt?')
+  end
+
+  #
+  # Suspends interaction with the interactive REPL interpreter
+  #
+  def _suspend
+    if (prompt_yesno('Background interactive SQL prompt?') == 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
+
+  def _winch
+    # noop
+  end
+
+  # 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
+
+    query
+  end
+
+  def _multiline
+    begin
+      require 'reline' unless defined?(::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 = ::Reline.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
+      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
+        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
+
+  attr_accessor :on_log_proc, :client_dispatcher
+
+end
+
+end
+end
+end
+end