Implement a method for command shells to register a post-session cleanup command

Author: busterb

lib/msf/base/sessions/command_shell.rb

@@ -47,10 +47,13 @@ def self.type
     "shell"
   end
 
-  def initialize(*args)
+  def initialize(conn, opts = {})
     self.platform ||= ""
     self.arch     ||= ""
     self.max_threads = 1
+    if !opts[:datastore]["CommandShellCleanupCommand"].blank?
+      @cleanup_command = opts[:datastore]["CommandShellCleanupCommand"]
+    end
     super
   end
 
@@ -193,12 +196,28 @@ def shell_write(buf)
   # :category: Msf::Session::Provider::SingleCommandShell implementors
   #
   # Closes the shell.
+  # Note: parent's 'self.kill' method calls cleanup below.
   #
   def shell_close()
-    rstream.close rescue nil
     self.kill
   end
 
+  ##
+  # :category: Msf::Session implementors
+  #
+  # Closes the shell.
+  #
+  def cleanup
+    if rstream
+      if !@cleanup_command.blank?
+        shell_command_token(@cleanup_command, 0)
+      end
+      rstream.close
+      rstream = nil
+    end
+    super
+  end
+
   #
   # Execute any specified auto-run scripts for this session
   #

lib/msf/base/sessions/command_shell_options.rb

@@ -17,9 +17,11 @@ def initialize(info = {})
 
     register_advanced_options(
       [
-        OptString.new('InitialAutoRunScript', [false, "An initial script to run on session creation (before AutoRunScript)", '']),
-        OptString.new('AutoRunScript', [false, "A script to run automatically on session creation.", ''])
-      ], self.class)
+        OptString.new('InitialAutoRunScript', "An initial script to run on session creation (before AutoRunScript)"),
+        OptString.new('AutoRunScript', "A script to run automatically on session creation."),
+        OptString.new('ShellCleanupCommand', "A command to run before the session is closed")
+      ]
+    )
   end
 
   def on_session(session)