From b9ec306046b8029098424dd609089dfae39173ed Mon Sep 17 00:00:00 2001
From: sjanusz-r7 <Simon_Janusz@rapid7.com>
Date: Wed, 3 Jan 2024 15:10:33 +0000
Subject: [PATCH] Use Reline for multi-line postgres query support in shell
 mode

---
 .../ui/console/command_dispatcher/db.rb       | 34 ++++++++++++-------
 1 file changed, 22 insertions(+), 12 deletions(-)

diff --git a/lib/rex/post/postgresql/ui/console/command_dispatcher/db.rb b/lib/rex/post/postgresql/ui/console/command_dispatcher/db.rb
index e7fb7dd8e5b9f..1892d33ebb713 100644
--- a/lib/rex/post/postgresql/ui/console/command_dispatcher/db.rb
+++ b/lib/rex/post/postgresql/ui/console/command_dispatcher/db.rb
@@ -1,7 +1,7 @@
 # -*- coding: binary -*-
 
 require 'pathname'
-require 'readline'
+require 'reline'
 
 module Rex
   module Post
@@ -97,21 +97,31 @@ def cmd_shell(*args)
               return
             end
 
-            multiline_query = false
-            raw_query = ''
-            while (line = ::Readline.readline("SQL #{'*' if multiline_query} > ", true))
-              raw_query << line << ' '
+            use_history = true
+            prompt_proc_before = ::Reline.prompt_proc
 
-              query_finished = !line.chomp.end_with?('\\')
-              break if query_finished
+            ::Reline.prompt_proc = proc { |line_buffer| line_buffer.each_with_index.map { |_line, i| i > 0 ? 'SQL *> ' : 'SQL >> ' } }
 
-              multiline_query = true
-            end
+            stop_words = %w[stop s exit e end quit q].freeze
+
+            finished = false
+            until finished
+              raw_query = ::Reline.readmultiline('SQL >> ', use_history) do |multiline_input|
+                if stop_words.include?(multiline_input.split.last)
+                  finished = true
+                  true
+                end
+                !multiline_input.split.last.end_with?('\\')
+              end
+              return if finished
 
-            # Format multi-line query
-            formatted_query = raw_query.split.map { |word| word unless word == '\\' }.join(' ')
+              formatted_query = raw_query.split.map { |word| word.chomp('\\') }.reject(&:empty?).compact.join(' ')
+
+              print_status "Running SQL Command: '#{formatted_query}'"
+              self.cmd_query(formatted_query)
+            end
 
-            self.cmd_query(formatted_query)
+            ::Reline.prompt_proc = prompt_proc_before
           end
 
           def cmd_query_help