Pull in common session functionality changes & fix PostgreSQL session…

… exit & code cleanup

lib/msf/base/sessions/postgresql.rb

@@ -15,8 +15,9 @@ class Msf::Sessions::PostgreSQL # < Msf::Sessions::CommandShell
   attr_accessor :platform
   attr_accessor :arch
 
-  # ##@param [PostgreSQL::Client] client
-  # @param [PostgreSQL::Client] rstream
+  # @param[Rex::IO::Stream] rstream
+  # @param [Hash] opts Options
+  # @param opts [PostgreSQL::Client] :client
   def initialize(rstream, opts={})
     @client = opts.fetch(:client)
     self.console = ::Rex::Post::PostgreSQL::Ui::Console.new(self)
@@ -35,7 +36,13 @@ def bootstrap(datastore = {}, handler = nil)
   end
 
   def process_autoruns(datastore)
-    # TODO - Implemented for now to keep things happy
+    ['InitialAutoRunScript', 'AutoRunScript'].each do |key|
+      next if datastore[key].nil? || datastore[key].empty?
+
+      args = Shellwords.shellwords(datastore[key])
+      print_status("Session ID #{session.sid} (#{session.tunnel_to_s}) processing #{key} '#{datastore[key]}'")
+      session.execute_script(args.shift, *args)
+    end
   end
 
   def type
@@ -62,6 +69,20 @@ def desc
     'PostgreSQL'
   end
 
+  def address
+    return @address if @address
+
+    @address, @port = self.client.conn.peerinfo.split(':')
+    @address
+  end
+
+  def port
+    return @port if @port
+
+    @address, @port = self.client.conn.peerinfo.split(':')
+    @port
+  end
+
   ##
   # :category: Msf::Session::Interactive implementors
   #
@@ -84,6 +105,10 @@ def reset_ui
     self.console.reset_ui
   end
 
+  def exit
+    self.console.stop
+  end
+
   protected
 
   ##

lib/rex/post/postgresql/ui.rb

@@ -1,3 +1,3 @@
 # -*- coding: binary -*-
 
-require 'rex/post/postgresql/ui/console'
+require 'rex/post/postgresql/ui/console'

lib/rex/post/postgresql/ui/console.rb

@@ -15,7 +15,7 @@ class Console
           # Dispatchers
           require 'rex/post/postgresql/ui/console/command_dispatcher'
           require 'rex/post/postgresql/ui/console/command_dispatcher/core'
-          require 'rex/post/postgresql/ui/console/command_dispatcher/db'
+          require 'rex/post/postgresql/ui/console/command_dispatcher/client'
 
           #
           # Initialize the PostgreSQL console.
@@ -36,7 +36,7 @@ def initialize(session)
             reset_ui
 
             enstack_dispatcher(::Rex::Post::PostgreSQL::Ui::Console::CommandDispatcher::Core)
-            enstack_dispatcher(::Rex::Post::PostgreSQL::Ui::Console::CommandDispatcher::DB)
+            enstack_dispatcher(::Rex::Post::PostgreSQL::Ui::Console::CommandDispatcher::Client)
 
             # Set up logging to whatever logsink 'core' is using
             if ! $dispatcher['postgresql']

lib/rex/post/postgresql/ui/console/command_dispatcher.rb

@@ -12,28 +12,7 @@ module Ui
         #
         ###
         module Console::CommandDispatcher
-          include Rex::Ui::Text::DispatcherShell::CommandDispatcher
-
-          #
-          # The hash of file names to class names after a module has already been
-          # loaded once on the client side.
-          #
-          # @@file_hash = {}
-          #
-          # #
-          # # Checks the file name to hash association to see if the module being
-          # # requested has already been loaded once.
-          # #
-          # def self.check_hash(name)
-          #   @@file_hash[name]
-          # end
-          #
-          # #
-          # # Sets the file path to class name association for future reference.
-          # #
-          # def self.set_hash(name, klass)
-          #   @@file_hash[name] = klass
-          # end
+          include Msf::Ui::Console::CommandDispatcher::Session
 
           #
           # Initializes an instance of the core command set using the supplied session and client
@@ -64,14 +43,6 @@ def session
             console.session
           end
 
-          def address
-            self.client.conn.remote_address.ip_address
-          end
-
-          def port
-            self.client.conn.remote_address.ip_port
-          end
-
           #
           # Returns the commands that meet the requirements
           #
@@ -86,7 +57,7 @@ def filter_commands(all, reqs)
 
           def unknown_command(cmd, line)
             if @filtered_commands.include?(cmd)
-              print_error("The \"#{cmd}\" command is not supported by this session type (#{client.session_type})")
+              print_error("The \"#{cmd}\" command is not supported by this session type (#{session.session_type})")
               return :handled
             end
 
@@ -98,7 +69,7 @@ def unknown_command(cmd, line)
           # to find usage documentation
           #
           def docs_dir
-            ::File.join(super, 'meterpreter')
+            ::File.join(super, 'postgresql_session')
           end
 
           #
@@ -110,7 +81,9 @@ def msf_loaded?
             return @msf_loaded unless @msf_loaded.nil?
 
             # if we get here we must not have initialized yet
-            @msf_loaded = !!(self.client.framework)
+
+            @msf_loaded = !session.framework.nil?
+            @msf_loaded
           end
 
           #
@@ -121,12 +94,10 @@ def log_error(msg)
 
             elog(msg, 'postgresql')
 
-            dlog("Call stack:\n#{$@.join("\n")}", 'postgresql')
+            dlog("Call stack:\n#{$ERROR_POSITION.join("\n")}", 'postgresql')
           end
-
         end
-
       end
     end
   end
-end
+end

lib/rex/post/postgresql/ui/console/command_dispatcher/db.rb → lib/rex/post/postgresql/ui/console/command_dispatcher/client.rb

@@ -10,10 +10,10 @@ module Ui
 
         ###
         #
-        # Core SMB client commands
+        # Core PostgreSQL client commands
         #
         ###
-        class Console::CommandDispatcher::DB
+        class Console::CommandDispatcher::Client
 
           include Rex::Post::PostgreSQL::Ui::Console::CommandDispatcher
 
@@ -37,13 +37,13 @@ class Console::CommandDispatcher::DB
             'daticurules' => 'ICU Rules',
             'datcollversion' => 'Collation Version',
             'datacl' => 'Access Privileges'
-          }
+          }.freeze
 
           #
           # Initializes an instance of the core command set using the supplied console
           # for interactivity.
           #
-          # @param [Rex::Post::SMB::Ui::Console] console
+          # @param [Rex::Post::PostgreSQL::Ui::Console] console
           def initialize(console)
             super
 
@@ -75,7 +75,7 @@ def commands
           end
 
           def name
-            'DB'
+            'PostgreSQL Client'
           end
 
           def help_args?(args)
@@ -87,7 +87,8 @@ def help_args?(args)
           def cmd_shell_help
             print_line 'Usage: shell'
             print_line
-            print_line 'Go into a raw SQL shell where SQL queries can be executed'
+            print_line 'Go into a raw SQL shell where SQL queries can be executed.'
+            print_line 'To exit, type `exit`, `quit`, `end` or `stop`.'
             print_line
           end
 
@@ -106,14 +107,22 @@ def cmd_shell(*args)
 
             finished = false
             until finished
-              raw_query = ::Reline.readmultiline('SQL >> ', use_history) do |multiline_input|
-                if stop_words.include?(multiline_input.split.last)
-                  finished = true
-                  true
+              begin
+                raw_query = ::Reline.readmultiline('SQL >> ', use_history) do |multiline_input|
+                  if stop_words.include?(multiline_input.split.last)
+                    finished = true
+                    true
+                  end
+                  !multiline_input.split.last.end_with?('\\')
                 end
-                !multiline_input.split.last.end_with?('\\')
+              rescue ::Interrupt
+                finished = true
+              end
+
+              if finished
+                print_status "Exiting Shell mode."
+                return
               end
-              return if finished
 
               formatted_query = raw_query.split.map { |word| word.chomp('\\') }.reject(&:empty?).compact.join(' ')