From 862e738015792986b59767350755bf2f2f35acdd Mon Sep 17 00:00:00 2001
From: Zach Goldman <zach_goldman@rapid7.com>
Date: Fri, 18 Aug 2023 10:45:34 -0500
Subject: [PATCH] extend error message for timeouts to include more detail to
 user

initial functionality, testing/cleanup still needed

script and command functionality

remove unnecessary accessor

switch puts to print_error in proc

ensure proc is reset, run on every error, add yard docs

fix yard, refactor/remove dead code

rename on_error_proc
---
 lib/msf/ui/console/command_dispatcher/core.rb | 55 ++++++++++++++-----
 lib/rex/post/meterpreter/ui/console.rb        | 27 +++++----
 lib/rex/ui/interactive.rb                     |  6 ++
 3 files changed, 64 insertions(+), 24 deletions(-)

diff --git a/lib/msf/ui/console/command_dispatcher/core.rb b/lib/msf/ui/console/command_dispatcher/core.rb
index 7eee0b163dc3..999d478bae21 100644
--- a/lib/msf/ui/console/command_dispatcher/core.rb
+++ b/lib/msf/ui/console/command_dispatcher/core.rb
@@ -1576,7 +1576,7 @@ def cmd_sessions(*args)
               rescue ::Rex::Post::Meterpreter::RequestError
                 print_error("Failed: #{$!.class} #{$!}")
               rescue Rex::TimeoutError
-                print_error("Operation timed out")
+                print_error("Operation timed out. Timeout currently #{session.response_timeout} seconds, you can configure this with %grnsessions -c <cmd> --timeout <value>%clr")
               end
             elsif session.type == 'shell' || session.type == 'powershell'
               output = session.shell_command(cmd)
@@ -1612,20 +1612,28 @@ def cmd_sessions(*args)
 
         cmds.each do |cmd|
           sessions.each do |session|
-            session = verify_session(session)
-            unless session.type == 'meterpreter'
-              print_error "Session ##{session.sid} is not a Meterpreter shell. Skipping..."
-              next
-            end
+            begin
+              session = verify_session(session)
+              unless session.type == 'meterpreter'
+                print_error "Session ##{session.sid} is not a Meterpreter shell. Skipping..."
+                next
+              end
 
-            next unless session
-            print_status("Running '#{cmd}' on #{session.type} session #{session.sid} (#{session.session_host})")
-            if session.respond_to?(:response_timeout)
-              last_known_timeout = session.response_timeout
-              session.response_timeout = response_timeout
-            end
+              next unless session
+              print_status("Running '#{cmd}' on #{session.type} session #{session.sid} (#{session.session_host})")
+              if session.respond_to?(:response_timeout)
+                last_known_timeout = session.response_timeout
+                session.response_timeout = response_timeout
+                session.on_run_command_error_proc = log_on_timeout_error("Send timed out. Timeout currently #{session.response_timeout} seconds, you can configure this with %grnsessions -C <cmd> --timeout <value>%clr")
+              end
 
-            output = session.run_cmd(cmd, driver.output)
+              output = session.run_cmd(cmd, driver.output)
+            ensure
+              if session.respond_to?(:response_timeout) && last_known_timeout
+                session.response_timeout = last_known_timeout
+                session.on_run_command_error_proc = nil
+              end
+            end
           end
         end
     when 'kill'
@@ -1674,16 +1682,19 @@ def cmd_sessions(*args)
           if session.respond_to?(:response_timeout)
             last_known_timeout = session.response_timeout
             session.response_timeout = response_timeout
+            session.on_run_command_error_proc = log_on_timeout_error("Send timed out. Timeout currently #{session.response_timeout} seconds, you can configure this with %grnsessions --interact <id> --timeout <value>%clr")
           end
           print_status("Starting interaction with #{session.name}...\n") unless quiet
           begin
             self.active_session = session
+
             sid = session.interact(driver.input.dup, driver.output)
             self.active_session = nil
             driver.input.reset_tab_completion if driver.input.supports_readline
           ensure
             if session.respond_to?(:response_timeout) && last_known_timeout
               session.response_timeout = last_known_timeout
+              session.on_run_command_error_proc = nil
             end
           end
         else
@@ -1708,12 +1719,14 @@ def cmd_sessions(*args)
           if session.respond_to?(:response_timeout)
             last_known_timeout = session.response_timeout
             session.response_timeout = response_timeout
+            session.on_run_command_error_proc = log_on_timeout_error("Send timed out. Timeout currently #{session.response_timeout} seconds, you can configure this with %grnsessions --timeout <value> --script <script> <id>%clr")
           end
           begin
             print_status("Session #{sess_id} (#{session.session_host}):")
             print_status("Running #{script} on #{session.type} session" +
                           " #{sess_id} (#{session.session_host})")
             begin
+              session.init_ui(driver.input, driver.output)
               session.execute_script(script, *extra)
             rescue ::Exception => e
               log_error("Error executing script or module: #{e.class} #{e}")
@@ -1721,7 +1734,9 @@ def cmd_sessions(*args)
           ensure
             if session.respond_to?(:response_timeout) && last_known_timeout
               session.response_timeout = last_known_timeout
+              session.on_run_command_error_proc = nil
             end
+            session.reset_ui
           end
         else
           print_error("Invalid session identifier: #{sess_id}")
@@ -2595,6 +2610,20 @@ def verify_session(session_id, quiet = false)
     end
   end
 
+  #
+  # Custom error code to handle timeout errors
+  #
+  # @param message [String] The message to be printed when a timeout error is hit
+  # @return [Proc] proc function that prints the specified error when the error types match
+  def log_on_timeout_error(message)
+    proc do |e|
+      next unless e.is_a?(Rex::TimeoutError) || e.is_a?(Timeout::Error)
+      elog(e)
+      print_error(message)
+      :handled
+    end
+  end
+
   #
   # Returns an array of lines at the provided line number plus any before and/or after lines requested
   # from all_lines by supplying the +before+ and/or +after+ parameters which are always positive
diff --git a/lib/rex/post/meterpreter/ui/console.rb b/lib/rex/post/meterpreter/ui/console.rb
index dccabf9782b4..f8c9d3fd927f 100644
--- a/lib/rex/post/meterpreter/ui/console.rb
+++ b/lib/rex/post/meterpreter/ui/console.rb
@@ -100,17 +100,22 @@ def queue_cmd(cmd)
   def run_command(dispatcher, method, arguments)
     begin
       super
-    rescue Timeout::Error
-      log_error("Operation timed out.")
-    rescue RequestError => info
-      log_error(info.to_s)
-    rescue Rex::InvalidDestination => e
-      log_error(e.message)
-    rescue ::Errno::EPIPE, ::OpenSSL::SSL::SSLError, ::IOError
-      self.client.kill
-    rescue  ::Exception => e
-      log_error("Error running command #{method}: #{e.class} #{e}")
-      elog(e)
+    rescue Exception => e
+      is_error_handled = self.client.on_run_command_error_proc && self.client.on_run_command_error_proc.call(e) == :handled
+      return if is_error_handled
+      case e
+      when Rex::TimeoutError, Rex::InvalidDestination
+        log_error(e.message)
+      when Timeout::Error
+        log_error('Operation timed out.')
+      when RequestError
+        log_error(e.to_s)
+      when ::Errno::EPIPE, ::OpenSSL::SSL::SSLError, ::IOError
+        self.client.kill
+      when ::Exception
+        log_error("Error running command #{method}: #{e.class} #{e}")
+        elog(e)
+      end
     end
   end
 
diff --git a/lib/rex/ui/interactive.rb b/lib/rex/ui/interactive.rb
index 5b5b3e22dc82..bb446409fe0c 100644
--- a/lib/rex/ui/interactive.rb
+++ b/lib/rex/ui/interactive.rb
@@ -128,6 +128,12 @@ def detach
   attr_accessor :on_print_proc
   attr_accessor :on_command_proc
 
+  #
+  # A function to be run when running a session command hits an error
+  #
+  # @return [Proc,nil] A function to be run when running a session command hits an error
+  attr_accessor :on_run_command_error_proc
+
 protected
 
   #