diff --git a/lib/rex/post/postgresql/ui/console.rb b/lib/rex/post/postgresql/ui/console.rb index b9fedda33664..9c45bac59dff 100644 --- a/lib/rex/post/postgresql/ui/console.rb +++ b/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 diff --git a/lib/rex/post/postgresql/ui/console/command_dispatcher/client.rb b/lib/rex/post/postgresql/ui/console/command_dispatcher/client.rb index 2196f33e6a5f..ced8dc4e8bfa 100644 --- a/lib/rex/post/postgresql/ui/console/command_dispatcher/client.rb +++ b/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 diff --git a/lib/rex/post/postgresql/ui/console/interactive_sql_client.rb b/lib/rex/post/postgresql/ui/console/interactive_sql_client.rb new file mode 100644 index 000000000000..28a450b8cb7b --- /dev/null +++ b/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