Skip to content

Commit

Permalink
SQL sessions have correct history manager support
Browse files Browse the repository at this point in the history
  • Loading branch information
sjanusz-r7 committed May 23, 2024
1 parent d37a825 commit 60c09e5
Show file tree
Hide file tree
Showing 4 changed files with 262 additions and 43 deletions.
57 changes: 53 additions & 4 deletions lib/msf/base/config.rb
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,8 @@ def self.get_config_root
'PluginDirectory' => "plugins",
'DataDirectory' => "data",
'LootDirectory' => "loot",
'LocalDirectory' => "local"
'LocalDirectory' => "local",
'HistoriesDirectory' => "histories"
}

##
Expand All @@ -97,6 +98,13 @@ def self.config_directory
self.new.config_directory
end

# Returns the histories directory default.
#
# @return [String] the SQL session histories directory.
def self.histories_directory
self.new.histories_directory
end

# Return the directory that logo files should be loaded from.
#
# @return [String] path to the logos directory.
Expand Down Expand Up @@ -235,20 +243,41 @@ def self.postgresql_session_history
self.new.postgresql_session_history
end

# Returns the full path to the PostgreSQL interactive query history file
#
# @return [String] path to the interactive query history file.
def self.postgresql_session_interactive_history
self.new.postgresql_session_interactive_history
end

# Returns the full path to the MSSQL session history file.
#
# @return [String] path to the history file.
def self.mssql_session_history
self.new.mssql_session_history
end

# Returns the full path to the MSSQL interactive query history file
#
# @return [String] path to the interactive query history file.
def self.mssql_session_interactive_history
self.new.mssql_session_interactive_history
end

# Returns the full path to the MySQL session history file.
#
# @return [String] path to the history file.
def self.mysql_session_history
self.new.mysql_session_history
end

# Returns the full path to the MySQL interactive query history file
#
# @return [String] path to the interactive query history file.
def self.mysql_session_interactive_history
self.new.mysql_session_interactive_history
end

def self.pry_history
self.new.pry_history
end
Expand Down Expand Up @@ -336,6 +365,13 @@ def config_directory
self['ConfigDirectory']
end

# Returns the histories directory default.
#
# @return [String] the SQL session histories directory.
def histories_directory
config_directory + FileSep + self['HistoriesDirectory']
end

# Returns the full path to the configuration file.
#
# @return [String] path to the configuration file.
Expand Down Expand Up @@ -363,15 +399,27 @@ def ldap_session_history
end

def postgresql_session_history
config_directory + FileSep + "postgresql_session_history"
histories_directory + FileSep + "postgresql_session_history"
end

def postgresql_session_interactive_history
histories_directory + FileSep + "postgresql_session_interactive_history"
end

def mysql_session_history
config_directory + FileSep + "mysql_session_history"
histories_directory + FileSep + "mysql_session_history"
end

def mysql_session_interactive_history
histories_directory + FileSep + "mysql_session_interactive_history"
end

def mssql_session_history
config_directory + FileSep + "mssql_session_history"
histories_directory + FileSep + "mssql_session_history"
end

def mssql_session_interactive_history
histories_directory + FileSep + "mssql_session_interactive_history"
end

def pry_history
Expand Down Expand Up @@ -495,6 +543,7 @@ def init
FileUtils.mkdir_p(user_module_directory)
FileUtils.mkdir_p(user_plugin_directory)
FileUtils.mkdir_p(user_data_directory)
FileUtils.mkdir_p(histories_directory)
end

# Loads configuration from the supplied file path, or the default one if
Expand Down
28 changes: 25 additions & 3 deletions lib/rex/post/sql/ui/console/interactive_sql_client.rb
Original file line number Diff line number Diff line change
Expand Up @@ -67,8 +67,20 @@ def _winch

# 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
name = session.type
query = {}
history_file = Msf::Config.send("#{name}_session_interactive_history")

# Multiline (Reline) and fallback (Readline) have separate history contexts as they are two different libraries.
framework.history_manager.with_context(history_file: history_file , name: name, input_library: :reline) do
query = _multiline
end

if query[:status] == :fail
framework.history_manager.with_context(history_file: history_file, name: name, input_library: :readline) do
query = _fallback
end
end

query
end
Expand Down Expand Up @@ -158,11 +170,21 @@ def _fallback
break if line.end_with? ';'
end

{ status: :success, result: line_buffer.join }
{ status: :success, result: line_buffer.join(' ') }
end

attr_accessor :on_log_proc, :client_dispatcher

private

def framework
client_dispatcher.shell.framework
end

def session
client_dispatcher.shell.session
end

end
end
end
Expand Down
81 changes: 60 additions & 21 deletions lib/rex/ui/text/shell/history_manager.rb
Original file line number Diff line number Diff line change
Expand Up @@ -24,10 +24,12 @@ def initialize
#
# @param [String,nil] history_file The file to load and persist commands to
# @param [String] name Human readable history context name
# @param [Symbol] input_library The input library to provide context for. :reline, :readline
# @param [Proc] block
# @return [nil]
def with_context(history_file: nil, name: nil, &block)
push_context(history_file: history_file, name: name)
def with_context(history_file: nil, name: nil, input_library: nil, &block)
# Default to Readline for backwards compatibility.
push_context(history_file: history_file, name: name, input_library: input_library || :readline)

begin
block.call
Expand Down Expand Up @@ -59,15 +61,21 @@ def _debug=(value)
@debug = value
end

# Return a queue of threads that have not yet finished saving the history file to disk.
# @return [Queue] a queue of threads that have not yet finished saving the history file to disk.
def _remaining_work
@remaining_work
end

private

def debug?
@debug
end

def push_context(history_file: nil, name: nil)
def push_context(history_file: nil, name: nil, input_library: nil)
$stderr.puts("Push context before\n#{JSON.pretty_generate(_contexts)}") if debug?
new_context = { history_file: history_file, name: name }
new_context = { history_file: history_file, name: name, input_library: input_library || :readline }

switch_context(new_context, @contexts.last)
@contexts.push(new_context)
Expand All @@ -91,47 +99,78 @@ def readline_available?
defined?(::Readline)
end

def reline_available?
begin
require 'reline'
defined?(::Reline)
rescue ::LoadError => _e
false
end
end

def clear_readline
return unless readline_available?

::Readline::HISTORY.length.times { ::Readline::HISTORY.pop }
end

def load_history_file(history_file)
return unless readline_available?
def clear_reline
return unless reline_available?

clear_readline
if File.exist?(history_file)
File.readlines(history_file).each do |e|
::Readline::HISTORY << e.chomp
end
::Reline::HISTORY.length.times { ::Reline::HISTORY.pop }
end

def load_history_file(context)
history_file = context[:history_file]
history = context[:input_library] == :reline ? ::Reline::HISTORY : ::Readline::HISTORY

begin
File.open(history_file, 'r') do |f|
context[:input_library] == :reline ? clear_reline : clear_readline
f.each do |line|
chomped_line = line.chomp
if context[:input_library] == :reline && history.last&.end_with?("\\")
history.last.delete_suffix!("\\")
history.last << "\n" << chomped_line
else
history << chomped_line
end
end
end
rescue Errno::EACCES, Errno::ENOENT => e
$stderr.puts("Failed to open history file: #{history_file} with error: #{e}") if debug?
end
end

def store_history_file(history_file)
return unless readline_available?
def store_history_file(context)
history_file = context[:history_file]
history = context[:input_library] == :reline ? ::Reline::HISTORY : ::Readline::HISTORY

history_diff = history.length < MAX_HISTORY ? history.length : MAX_HISTORY

cmds = []
history_diff = ::Readline::HISTORY.length < MAX_HISTORY ? ::Readline::HISTORY.length : MAX_HISTORY
history_diff.times do
entry = ::Readline::HISTORY.pop
cmds.push(entry) unless entry.nil?
entry = history.pop
cmds << entry.scrub.split("\n").join("\\\n")
end

write_history_file(history_file, cmds)
write_history_file(history_file, cmds.reverse)
end

def switch_context(new_context, old_context=nil)
if old_context && old_context[:history_file]
store_history_file(old_context[:history_file])
store_history_file(old_context)
end

if new_context && new_context[:history_file]
load_history_file(new_context[:history_file])
load_history_file(new_context)
else
clear_readline
clear_reline
end
rescue SignalException => e
rescue SignalException => _e
clear_readline
clear_reline
end

def write_history_file(history_file, cmds)
Expand All @@ -144,7 +183,7 @@ def write_history_file(history_file, cmds)
cmds = event[:cmds]

File.open(history_file, 'wb+') do |f|
f.puts(cmds.reverse)
f.puts(cmds)
end

rescue => e
Expand Down
Loading

0 comments on commit 60c09e5

Please sign in to comment.