From af60ee2f19e3b3a61b3edbc375857c3779d08bfd Mon Sep 17 00:00:00 2001
From: Zach Goldman <zach_goldman@rapid7.com>
Date: Wed, 18 Oct 2023 16:38:06 -0500
Subject: [PATCH] ensure proc is reset, run on every error, add yard docs

---
 lib/msf/ui/console/command_dispatcher/core.rb | 68 ++++++++++++-------
 lib/rex/post/meterpreter/ui/console.rb        | 32 ++++-----
 lib/rex/ui/interactive.rb                     |  5 ++
 3 files changed, 63 insertions(+), 42 deletions(-)

diff --git a/lib/msf/ui/console/command_dispatcher/core.rb b/lib/msf/ui/console/command_dispatcher/core.rb
index 67c03702d73f4..9e543c75097d8 100644
--- a/lib/msf/ui/console/command_dispatcher/core.rb
+++ b/lib/msf/ui/console/command_dispatcher/core.rb
@@ -957,14 +957,6 @@ def cmd_load_help
     print_line
   end
 
-  def log_on_timeout_error(message)
-    proc do |e|
-      next unless e.is_a?(Rex::TimeoutError)
-      elog(e)
-      print_error(message)
-    end
-  end
-
   def list_plugins
     plugin_directories = {
       'Framework' => Msf::Config.plugin_directory,
@@ -1584,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)
@@ -1620,21 +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
-              session.on_error_proc = log_on_timeout_error("Send timed out. Timeout currently #{session.response_timeout} seconds, you can configure this with sessions -C <cmd> --timeout <value>")
-            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_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_error_proc = nil
+              end
+            end
           end
         end
     when 'kill'
@@ -1683,19 +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_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
 
-            session.on_error_proc = log_on_timeout_error("Send timed out. Timeout currently #{session.response_timeout} seconds, you can configure this with sessions --interact <id> --timeout <value>")
-
             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_error_proc = nil
             end
           end
         else
@@ -1720,13 +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_error_proc = log_on_timeout_error("Send timed out. Timeout currently #{session.response_timeout} seconds, you can configure this with sessions --timeout <value> --script <script> <id>")
+            session.on_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}")
@@ -1734,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_error_proc = nil
             end
+            session.reset_ui
           end
         else
           print_error("Invalid session identifier: #{sess_id}")
@@ -2608,6 +2610,22 @@ 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
+
+  def log_on_timeout_error(message)
+    proc do |e|
+      handled = false
+      if e.is_a?(Rex::TimeoutError) || e.is_a?(Timeout::Error)
+        elog(e)
+        print_error(message)
+        :handled
+      end
+    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 417e60203a8fd..3370499f25e96 100644
--- a/lib/rex/post/meterpreter/ui/console.rb
+++ b/lib/rex/post/meterpreter/ui/console.rb
@@ -99,24 +99,22 @@ def queue_cmd(cmd)
   #
   def run_command(dispatcher, method, arguments)
     begin
-      super 
-    rescue Rex::TimeoutError => e
-      if self.client.on_error_proc
-        self.client.on_error_proc.call(e)
-      else
-        log_error(e.message)
+      super
+    rescue Exception => e
+      unless self.client.on_error_proc && self.client.on_error_proc.call(e) == :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}")
+        end
       end
-    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)
     end
   end
 
diff --git a/lib/rex/ui/interactive.rb b/lib/rex/ui/interactive.rb
index 6819e21496d23..646d344a3c932 100644
--- a/lib/rex/ui/interactive.rb
+++ b/lib/rex/ui/interactive.rb
@@ -127,6 +127,11 @@ def detach
 
   attr_accessor :on_print_proc
   attr_accessor :on_command_proc
+
+  #
+  # A function to be run when the session hits an error
+  #
+
   attr_accessor :on_error_proc
 
 protected