lib/rex/post/sql/ui/console/command_dispatcher/client.rb
# -*- coding: binary -*-
module Rex
module Post
module Sql
module Ui
module Console
module CommandDispatcher
###
#
# Core Generic SQL client commands
#
###
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.
#
# @param console The protocol-specific Rex Ui Console
def initialize(console)
super
@db_search_results = []
end
#
# List of supported commands.
#
def commands
cmds = {
'query' => 'Run a single SQL query',
'query_interactive' => 'Enter an interactive prompt for running multiple SQL queries',
}
reqs = {}
filter_commands(cmds, reqs)
end
# @return [String] The name of the client
def name
raise ::NotImplementedError
end
# @param [Array] args An array of arguments passed in to a command
# @return [TrueClass, FalseClass] True if the array contains '-h' or '--help', else false.
def help_args?(args)
return false unless args.instance_of?(::Array)
args.include?('-h') || args.include?('--help')
end
def cmd_query_interactive_help
print_line 'Usage: query_interactive'
print_line
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 query_interactive_help
print_line 'Interactive SQL prompt'
print_line
print_line 'You are in an interactive SQL shell where SQL queries can be executed.'
print_line 'SQL commands ending with ; will be executed on the remote server.'
print_line "To exit, type 'exit', 'quit', 'end' or 'stop'."
print_line
end
def cmd_query_interactive(*args)
if help_args?(args)
cmd_query_interactive_help
return
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 normalise_sql_result(result)
raise ::NotImplementedError
end
# Take in a normalised SQL result and print it.
# If there are any errors, print those instead.
# @param [Hash {Symbol => Array, Symbol => Array, Symbol => Array}] result A hash of 'rows', 'columns' and 'errors'
def format_result(result)
if result[:errors].any?
return "Query has failed. Reasons: #{result[:errors].join(', ')}"
end
number_column = ['#']
columns = [number_column, result[:columns]].flatten
rows = result[:rows].map.each_with_index do |row, i|
[i, row].flatten
end
::Rex::Text::Table.new(
'Header' => 'Response',
'Indent' => 4,
'Columns' => columns,
'Rows' => rows
)
end
def cmd_query_help
raise ::NotImplementedError
end
def run_query(query)
begin
result = client.query(query)
rescue => e
elog("Running query '#{query}' failed on session #{self.inspect}", error: e)
return { status: :error, result: { errors: [e] } }
end
if result.respond_to?(:cmd_tag) && result.cmd_tag
print_status result.cmd_tag
print_line
end
{ status: :success, result: result }
end
def cmd_query(*args)
show_help = false
call_interactive = false
@@query_opts.parse(args) do |opt, _idx, val|
case opt
when nil
show_help = true if val == 'help'
break
when '-h', '--help'
show_help = true
break
when '-i', '--interact'
call_interactive = true
break
end
end
if args.empty? || show_help
cmd_query_help
return
end
if call_interactive
cmd_query_interactive
return
end
result = run_query(args.join(' '))
case result[:status]
when :success
# When changing a database in MySQL, we get a nil result back.
if result[:result].nil?
print_status 'Query executed successfully'
return
end
normalised_result = normalise_sql_result(result[:result])
# MSSQL returns :success, even if the query failed due to wrong syntax.
if normalised_result[:errors].any?
print_error "Query has failed. Reasons: #{normalised_result[:errors].join(', ')}"
return
end
# When changing a database in MSSQL, we get a result, but it doesn't contain colnames or rows.
if normalised_result[:rows].nil? || normalised_result[:columns].nil?
print_status 'Query executed successfully'
return
end
formatted_result = format_result(normalised_result)
print_line(formatted_result.to_s)
when :error
print_error "Query has failed. Reasons: #{result[:result][:errors].join(', ')}"
result[:result][:errors].each do |error|
handle_error(error)
end
else
elog "Unknown query status: #{result[:status]}"
print_error "Unknown query status: #{result[:status]}"
end
end
def process_query(query: '')
return '' if query.empty?
query.lines.each.map { |line| line.chomp.chomp('\\').strip }.reject(&:empty?).compact.join(' ')
end
# Handles special cased error for each protocol
# that are return when a session has died resulting in a session
# needing to be closed
#
# @param [Object] e
def handle_error(e)
case e
when EOFError
_close_session
end
end
private
# Sets session.alive to close sessions and handle setting session.client.interacting
# if currently in the context of the `query_interactive` command
#
# @return [FalseClass]
def _close_session
session.alive = false
session.client.interacting = false if session.client && session.client.respond_to?(:interacting)
end
end
end
end
end
end
end
end