Skip to content

Commit

Permalink
addresses feedback on code structure and search behavior
Browse files Browse the repository at this point in the history
  • Loading branch information
@zgoldman-r7
zgoldman-r7 committed Oct 10, 2023
1 parent 9683743 commit 6642afd
Show file tree
Hide file tree
Showing 3 changed files with 160 additions and 114 deletions.
17 changes: 5 additions & 12 deletions lib/msf/base/serializer/readable_text.rb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -834,7 +834,7 @@ def self.dump_datastore(name, ds, indent = DefaultIndent, col = DefaultColumnWra
def self.dump_sessions(framework, opts={})
output = ""
verbose = opts[:verbose] || false
session_ids = opts[:session_ids] || nil
sessions = opts[:sessions] || framework.sessions
show_active = opts[:show_active] || false
show_inactive = opts[:show_inactive] || false
# if show_active and show_inactive are false the caller didn't
Expand All @@ -861,12 +861,10 @@ def self.dump_sessions(framework, opts={})
'Columns' => columns,
'Indent' => indent)

framework.sessions.each { |k|
next unless session_ids.nil? || session_ids.include?(k[0])
session = k[1]
sessions.each do |session_id, session|
row = create_msf_session_row(session, show_extended)
tbl << row
}
end

output << (tbl.rows.count > 0 ? tbl.to_s : "#{tbl.header_to_s}No active sessions.\n")
end
Expand Down Expand Up @@ -987,14 +985,9 @@ def self.dump_sessions_verbose(framework, opts={})
return out
end

framework.sessions.each_sorted do |k|
session = framework.sessions[k]
sessions = opts[:sessions] || framework.sessions

if opts[:session_ids]
unless opts[:session_ids].include? session.sid
next
end
end
sessions.each do |session_id, session|
sess_info = session.info.to_s
sess_id = session.sid.to_s
sess_name = session.sname.to_s
Expand Down
146 changes: 87 additions & 59 deletions lib/msf/ui/console/command_dispatcher/core.rb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -114,12 +114,14 @@ class Core
["-g", "--global"] => [ false, "Operate on global datastore variables"]
)

VALID_PARAMS =
VALID_SESSION_SEARCH_PARAMS =
%w[
last_checkin
session_id
session_type
]

private_constant :VALID_SESSION_SEARCH_PARAMS

# Returns the list of commands supported by this command dispatcher
def commands
Expand Down Expand Up @@ -1431,7 +1433,6 @@ def cmd_sessions(*args)
begin
method = nil
quiet = false
search = false
show_active = false
show_inactive = false
show_extended = false
Expand Down Expand Up @@ -1536,6 +1537,14 @@ def cmd_sessions(*args)
print_warning("Database not connected; list of inactive sessions unavailable")
end

if search_term
matching_sessions = get_matching_sessions(search_term)
if matching_sessions.empty?
print_status("No matching sessions.")
return
end
end

last_known_timeout = nil

# Now, perform the actual method
Expand All @@ -1545,13 +1554,7 @@ def cmd_sessions(*args)
print_error("No command specified!")
return false
end
if search_term
matching_sessions = get_matching_sessions(search_term)
if matching_sessions.empty?
print_error("No matching sessions.")
return
end
end

cmds.each do |cmd|
if sid
sessions = session_list
Expand Down Expand Up @@ -1666,25 +1669,16 @@ def cmd_sessions(*args)
end
end
when 'killall'
if search_term
matching_sessions = get_matching_sessions(search_term)
if matching_sessions.empty?
print_status("No matching sessions.")
return
end
if matching_sessions
print_status("Killing matching sessions...")
print_line
print(Serializer::ReadableText.dump_sessions(framework, show_active: show_active, show_inactive: show_inactive, show_extended: show_extended, verbose: verbose, session_ids: matching_sessions.keys()))
print(Serializer::ReadableText.dump_sessions(framework, show_active: show_active, show_inactive: show_inactive, show_extended: show_extended, verbose: verbose, sessions: matching_sessions))
print_line
else
matching_sessions = framework.sessions
print_status("Killing all sessions...")
end
matching_sessions.each do |session_id, session|
unless matching_sessions.nil? || matching_sessions.include?(session_id)
next
end

if session
if session.respond_to?(:response_timeout)
last_known_timeout = session.response_timeout
Expand Down Expand Up @@ -1786,18 +1780,8 @@ def cmd_sessions(*args)
end
end
when 'list', 'list_inactive', nil
if search_term
matching_sessions = get_matching_sessions(search_term)
if matching_sessions.empty?
print_error("No matching sessions.")
return
end
end
if matching_sessions
session_ids = matching_sessions.keys()
end
print_line
print(Serializer::ReadableText.dump_sessions(framework, show_active: show_active, show_inactive: show_inactive, show_extended: show_extended, verbose: verbose, session_ids: session_ids))
print(Serializer::ReadableText.dump_sessions(framework, show_active: show_active, show_inactive: show_inactive, show_extended: show_extended, verbose: verbose, sessions: matching_sessions))
print_line
when 'name'
if session_name.blank?
Expand Down Expand Up @@ -1839,79 +1823,123 @@ def cmd_sessions(*args)
def get_matching_sessions(search_term)
matching_sessions = {}
terms = search_term.split
id_searches = []
type_searches = []
checkin_searches = []
terms.each do |term|
matches = filter_sessions_by_search(term)
matches.each do |session_id, session|
matching_sessions.store(session_id, session)
case term.split(":").first
when "session_id"
id_searches << term
when "session_type"
type_searches << term
when "last_checkin"
checkin_searches << term
else
raise ArgumentError, "Please provide valid search term. Given: #{term.split(":").first}"
end
end
searches = []
unless id_searches.empty?
id_matches = {}
id_searches.each do |term|
matches = filter_sessions_by_search(term)
id_matches = id_matches.merge(matches)
end
searches << id_matches
end
unless type_searches.empty?
type_matches = {}
type_searches.each do |term|
matches = filter_sessions_by_search(term)
type_matches = type_matches.merge(matches)
end
searches << type_matches
end
unless checkin_searches.empty?
checkin_matches = {}
first_loop = true
checkin_searches.each do |term|
matches = filter_sessions_by_search(term)
if first_loop
checkin_matches = matches
else
checkin_matches = checkin_matches.select{ |session_id, session| matches[session_id] == session }
end
end
searches << checkin_matches
end
if searches
matching_sessions = searches.first
searches[1..].each do |result_set|
matching_sessions = matching_sessions.select{| session_id, session| result_set[session_id] == session}
end
else
raise ArgumentError, "Please provide a valid search query."
end
matching_sessions
end

def filter_sessions_by_search(search_term)
matching_sessions = {}

field = search_term.split(":")[0]
framework.sessions.each do |session_id, session|
case search_term.split(":")[0]
case field
when "last_checkin"
if session.respond_to?(:last_checkin) && session.last_checkin && evaluate_search_criteria(session, search_term)
matching_sessions.store(session_id, session)
matching_sessions[session_id] = session
end
when "session_type"
matching_sessions.store(session_id, session) if evaluate_search_criteria(session, search_term)
when "session_id"
matching_sessions.store(session_id, session) if evaluate_search_criteria(session, search_term)
when "session_type", "session_id"
matching_sessions[session_id] = session if evaluate_search_criteria(session, search_term)
else
raise ArgumentError, "Unrecognized search term: #{field}"
end
end
matching_sessions
end

def evaluate_search_criteria(session, search_term)
parts = search_term.split(":")
field, operator, value = search_term.split(":")

field = parts[0]
operator = parts[1]
value = parts[2] if parts[2]
return false unless VALID_PARAMS.include? parts[0]
raise ArgumentError, "Unrecognized search term: #{field}" unless VALID_SESSION_SEARCH_PARAMS.include? field

case field
when "last_checkin"
checkin_time = session.last_checkin
threshold_time = Time.now - parse_duration(value)
case operator
when "before"
return checkin_time > threshold_time
when "after"
return checkin_time < threshold_time
when "after"
return checkin_time > threshold_time
else
return false
end
when "session_id"
return session.sid.to_s == operator
when "session_type"
return session.type == operator
return session.type.casecmp?(operator)
end
end

def parse_duration(duration)
total_time = 0
time_tokens = duration.scan(/\d+/).zip(duration.scan(/[a-zA-Z]+/))
time_tokens.each do |pair|
raise "Please specify both time units and amounts" if pair[1].nil?
case pair[1]
time_tokens = duration.scan(/(?:\d+\.?\d*|\.\d+)/).zip(duration.scan(/[a-zA-Z]+/))
time_tokens.each do |value, unit|
raise ArgumentError, "Please specify both time units and amounts" if unit.nil? || value.nil?
case unit
when "d"
total_time = total_time + pair[0].to_i * 86400
total_time += value.to_f * 86400
when "h"
total_time = total_time + pair[0].to_i * 3600
total_time += value.to_f * 3600
when "m"
total_time = total_time + pair[0].to_i * 60
total_time += value.to_f * 60
when "s"
total_time = total_time + pair[0].to_i
total_time += value.to_f
else
raise "Unrecognized time format"
raise ArgumentError, "Unrecognized time format: #{value}"
end
end
return total_time
total_time.to_i
end

#
Expand Down
Loading

0 comments on commit 6642afd

Please sign in to comment.