Skip to content

Commit

Permalink
Use Reline for multi-line postgres query support in shell mode
Browse files Browse the repository at this point in the history
  • Loading branch information
sjanusz-r7 committed Jan 3, 2024
1 parent 3b2031c commit b9ec306
Showing 1 changed file with 22 additions and 12 deletions.
34 changes: 22 additions & 12 deletions lib/rex/post/postgresql/ui/console/command_dispatcher/db.rb
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
# -*- coding: binary -*-

require 'pathname'
require 'readline'
require 'reline'

module Rex
module Post
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit b9ec306

Please sign in to comment.