rename shell to query_interactive for sql session types, add -i flag

zgoldman-r7 · Feb 27, 2024 · afbd18f · afbd18f
1 parent b91430c
commit afbd18f
Show file tree

Hide file tree

Showing 6 changed files with 32 additions and 19 deletions.
diff --git a/lib/rex/post/mssql/ui/console/command_dispatcher/client.rb b/lib/rex/post/mssql/ui/console/command_dispatcher/client.rb
@@ -24,8 +24,8 @@ def name
           def cmd_query_help
             print_line 'Usage: query'
             print_line
-            print_line 'Run a raw SQL query on the target.'
-            print_line
+            print_line 'Run a single SQL query on the target.'
+            print_line @@query_opts.usage
             print_line 'Examples:'
             print_line
             print_line '    query select @@version;'

diff --git a/lib/rex/post/mysql/ui/console/command_dispatcher/client.rb b/lib/rex/post/mysql/ui/console/command_dispatcher/client.rb
@@ -21,8 +21,8 @@ def name
           def cmd_query_help
             print_line 'Usage: query'
             print_line
-            print_line 'Run a raw SQL query on the target.'
-            print_line
+            print_line 'Run a single SQL query on the target.'
+            print_line @@query_opts.usage
             print_line 'Examples:'
             print_line
             print_line '    query SHOW DATABASES;'

diff --git a/lib/rex/post/postgresql/ui/console/command_dispatcher/client.rb b/lib/rex/post/postgresql/ui/console/command_dispatcher/client.rb
@@ -25,8 +25,8 @@ def name
           def cmd_query_help
             print_line 'Usage: query'
             print_line
-            print_line 'Run a raw SQL query on the target.'
-            print_line
+            print_line 'Run a single SQL query on the target.'
+            print_line @@query_opts.usage
             print_line 'Examples:'
             print_line
             print_line '    query SELECT user;'

diff --git a/lib/rex/post/sql/ui/console/command_dispatcher/client.rb b/lib/rex/post/sql/ui/console/command_dispatcher/client.rb
@@ -16,6 +16,11 @@ module Client
 
               include Rex::Post::Sql::Ui::Console::CommandDispatcher
 
+              @@query_opts = Rex::Parser::Arguments.new(
+                ['-h', '--help'] => [false, 'Help menu.'],
+                ['-i', '--interact'] => [false,  'Enter an interactive prompt for running multiple SQL queries'],
+                )
+
               #
               # Initializes an instance of the core command set using the supplied console
               # for interactivity.
@@ -32,8 +37,8 @@ def initialize(console)
               #
               def commands
                 cmds = {
-                  'query'   => 'Run a raw SQL query',
-                  'shell'   => 'Enter a raw shell where SQL queries can be executed',
+                  'query'   => 'Run a single SQL query',
+                  'query_interactive'   => 'Enter an interactive prompt for running multiple SQL queries',
                 }
 
                 reqs = {}
@@ -54,17 +59,17 @@ def help_args?(args)
                 args.include?('-h') || args.include?('--help')
               end
 
-              def cmd_shell_help
-                print_line 'Usage: shell'
+              def cmd_query_interactive_help
+                print_line 'Usage: query_interactive'
                 print_line
-                print_line 'Go into a raw SQL shell where SQL queries can be executed.'
+                print_line 'Go into an interactive SQL shell where SQL queries can be executed.'
                 print_line "To exit, type 'exit', 'quit', 'end' or 'stop'."
                 print_line
               end
 
-              def cmd_shell(*args)
+              def cmd_query_interactive(*args)
                 if help_args?(args)
-                  cmd_shell_help
+                  cmd_query_interactive_help
                   return
                 end
 
@@ -120,9 +125,15 @@ def run_query(query)
               end
 
               def cmd_query(*args)
-                if help_args?(args)
-                  cmd_query_help
-                  return
+                @@query_opts.parse(args) do |opt, idx, val|
+                  case opt
+                  when '-h', '--help'
+                    cmd_query_help
+                    return
+                  when '-i', '--interact'
+                    cmd_query_interactive
+                    return
+                  end
                 end
 
                 result = run_query(args.join(' '))

diff --git a/lib/rex/post/sql/ui/console/interactive_sql_client.rb b/lib/rex/post/sql/ui/console/interactive_sql_client.rb
@@ -101,7 +101,7 @@ def _multiline
 
     if finished
       self.interacting = false
-      print_status 'Exiting Shell mode.'
+      print_status 'Exiting Interactive mode.'
       return { status: :exit, result: nil }
     end
 
@@ -116,7 +116,7 @@ def _fallback
 
       if stop_words.include? line.chomp.downcase
         self.interacting = false
-        print_status 'Exiting Shell mode.'
+        print_status 'Exiting Interactive mode.'
         return { status: :exit, result: nil }
       end
 

diff --git a/lib/rex/ui/text/dispatcher_shell.rb b/lib/rex/ui/text/dispatcher_shell.rb
@@ -237,6 +237,8 @@ def help_to_s(opts={})
       # If this dispatcher has no commands, we can't do anything useful.
       return "" if commands.nil? or commands.length == 0
 
+      # Set increase width of commands column if values longer than 12 characters exist, but cap at 20
+      command_width = [[12, commands.keys.max_by(&:length).length].max, 20].min
       # Display the commands
       tbl = Rex::Text::Table.new(
         'Header'  => "#{self.name} Commands",
@@ -250,7 +252,7 @@ def help_to_s(opts={})
           {
             'Command' =>
               {
-                'Width' => 12
+                'Width' => command_width
               }
           })