Skip to content

Commit

Permalink
Land #18933, update SQL sessions to correctly manage history
Browse files Browse the repository at this point in the history
  • Loading branch information
adfoster-r7 authored Mar 28, 2024
2 parents 37d3c88 + e2814d6 commit c0d66fd
Show file tree
Hide file tree
Showing 4 changed files with 117 additions and 39 deletions.
33 changes: 33 additions & 0 deletions lib/msf/base/config.rb
Original file line number Diff line number Diff line change
Expand Up @@ -228,20 +228,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_history_interactive
self.new.postgresql_session_history_interactive
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_history_interactive
self.new.mssql_session_history_interactive
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_history_interactive
self.new.mysql_session_history_interactive
end

def self.pry_history
self.new.pry_history
end
Expand Down Expand Up @@ -355,14 +376,26 @@ def postgresql_session_history
config_directory + FileSep + "postgresql_session_history"
end

def postgresql_session_history_interactive
postgresql_session_history + "_interactive"
end

def mysql_session_history
config_directory + FileSep + "mysql_session_history"
end

def mysql_session_history_interactive
mysql_session_history + "_interactive"
end

def mssql_session_history
config_directory + FileSep + "mssql_session_history"
end

def mssql_session_history_interactive
mssql_session_history + "_interactive"
end

def pry_history
config_directory + FileSep + "pry_history"
end
Expand Down
25 changes: 23 additions & 2 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,19 @@ 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 = {}

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

if query[:status] == :fail
framework.history_manager.with_context(history_file: Msf::Config.send("#{name}_session_history_interactive"), name: name, input_library: :readline) do
query = _fallback
end
end

query
end
Expand Down Expand Up @@ -163,6 +174,16 @@ def _fallback

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
68 changes: 46 additions & 22 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 @@ -65,9 +67,9 @@ 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 +93,69 @@ 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?

::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

clear_readline
if File.exist?(history_file)
File.readlines(history_file).each do |e|
::Readline::HISTORY << e.chomp
File.open(history_file, 'r') do |f|
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
end
end

def store_history_file(history_file)
return unless readline_available?
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?
end
def store_history_file(context)
history_file = context[:history_file]
history = context[:input_library] == :reline ? ::Reline::HISTORY : ::Readline::HISTORY

history_to_save = history.map { |line| line.scrub.split("\n").join("\\\n") }

write_history_file(history_file, cmds)
write_history_file(history_file, history_to_save)
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 +168,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
30 changes: 15 additions & 15 deletions spec/lib/rex/ui/text/shell/history_manager_spec.rb
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@
(expect do |block|
subject.with_context(name: 'a') do
expected_contexts = [
{ history_file: nil, name: 'a' },
{ history_file: nil, input_library: :readline, name: 'a' },
]
expect(subject._contexts).to eq(expected_contexts)
block.to_proc.call
Expand All @@ -36,15 +36,15 @@

context 'when there is an existing stack' do
before(:each) do
subject.send(:push_context, history_file: nil, name: 'a')
subject.send(:push_context, history_file: nil, input_library: :readline, name: 'a')
end

it 'continues to have the previous existing stack' do
subject.with_context {
# noop
}
expected_contexts = [
{ history_file: nil, name: 'a' },
{ history_file: nil, input_library: :readline, name: 'a' },
]
expect(subject._contexts).to eq(expected_contexts)
end
Expand All @@ -53,8 +53,8 @@
(expect do |block|
subject.with_context(name: 'b') do
expected_contexts = [
{ history_file: nil, name: 'a' },
{ history_file: nil, name: 'b' },
{ history_file: nil, input_library: :readline, name: 'a' },
{ history_file: nil, input_library: :readline, name: 'b' },
]
expect(subject._contexts).to eq(expected_contexts)
block.to_proc.call
Expand All @@ -69,7 +69,7 @@
}
end.to raise_exception ArgumentError, 'Mock error'
expected_contexts = [
{ history_file: nil, name: 'a' },
{ history_file: nil, input_library: :readline, name: 'a' },
]
expect(subject._contexts).to eq(expected_contexts)
end
Expand All @@ -79,9 +79,9 @@
describe '#push_context' do
context 'when the stack is empty' do
it 'stores the history contexts' do
subject.send(:push_context, history_file: nil, name: 'a')
subject.send(:push_context, history_file: nil, input_library: :readline, name: 'a')
expected_contexts = [
{ history_file: nil, name: 'a' }
{ history_file: nil, input_library: :readline, name: 'a' }
]
expect(subject._contexts).to eq(expected_contexts)
end
Expand All @@ -90,12 +90,12 @@
context 'when multiple values are pushed' do
it 'stores the history contexts' do
subject.send(:push_context, history_file: nil, name: 'a')
subject.send(:push_context, history_file: nil, name: 'b')
subject.send(:push_context, history_file: nil, name: 'c')
subject.send(:push_context, history_file: nil, input_library: :readline, name: 'b')
subject.send(:push_context, history_file: nil, input_library: :reline, name: 'c')
expected_contexts = [
{ history_file: nil, name: 'a' },
{ history_file: nil, name: 'b' },
{ history_file: nil, name: 'c' },
{ history_file: nil, input_library: :readline, name: 'a' },
{ history_file: nil, input_library: :readline, name: 'b' },
{ history_file: nil, input_library: :reline, name: 'c' },
]
expect(subject._contexts).to eq(expected_contexts)
end
Expand All @@ -113,12 +113,12 @@
end

context 'when the stack is not empty' do
it 'continues to have an empty stack' do
it 'continues to have a non-empty stack' do
subject.send(:push_context, history_file: nil, name: 'a')
subject.send(:push_context, history_file: nil, name: 'b')
subject.send(:pop_context)
expected_contexts = [
{ history_file: nil, name: 'a' },
{ history_file: nil, input_library: :readline, name: 'a' },
]
expect(subject._contexts).to eq(expected_contexts)
end
Expand Down

0 comments on commit c0d66fd

Please sign in to comment.