PostgreSQL session improvements

lib/metasploit/framework/login_scanner/postgres.rb

@@ -11,6 +11,14 @@ module LoginScanner
       class Postgres
         include Metasploit::Framework::LoginScanner::Base
 
+        # @returns [Boolean] If a login is successful and this attribute is true - a PostgreSQL::Client instance is used as proof,
+        #   and the socket is not immediately closed
+        attr_accessor :use_client_as_proof
+
+        # @!attribute dispatcher
+        #   @return [PostgreSQL::Dispatcher::Socket]
+        attr_accessor :dispatcher
+
         DEFAULT_PORT         = 5432
         DEFAULT_REALM        = 'template1'
         LIKELY_PORTS         = [ DEFAULT_PORT ]
@@ -70,8 +78,16 @@ def attempt_login(credential)
           end
 
           if pg_conn
-            pg_conn.close
             result_options[:status] = Metasploit::Model::Login::Status::SUCCESSFUL
+
+            # This module no long owns the socket, return it as proof so the calling context can perform additional operations
+            # Additionally assign values to nil to avoid closing the socket etc automatically
+            if use_client_as_proof
+              result_options[:proof] = pg_conn
+              pg_conn = nil
+            else
+              pg_conn.close
+            end
           else
             result_options[:status] = Metasploit::Model::Login::Status::INCORRECT
           end

lib/msf/base/sessions/postgresql.rb

@@ -10,20 +10,16 @@ class Msf::Sessions::PostgreSQL # < Msf::Sessions::CommandShell
 
   # @return [Rex::Post::PostgreSQL::Ui::Console] The interactive console
   attr_accessor :console
-  # PostgreSQL::Client maybe?
-  # @return [RubySMB::Client] The SMB client
   # @return [PostgreSQL::Client]
   attr_accessor :client
   attr_accessor :platform
   attr_accessor :arch
 
-  # TODO: Confirm existing terminology, smb_client vs client vs dispatcher vs something else
-  # @param [RubySMB::Client] client
+  # ##@param [PostgreSQL::Client] client
+  # @param [PostgreSQL::Client] rstream
   def initialize(rstream, opts={})
     @client = opts.fetch(:client)
-    # @info = "SMB TODO"
-    #self.console = Rex::Post::SMB::Ui::Console.new(self)
-    self.console = Rex::Post::PostgreSQL::Ui::Console.new(self)
+    self.console = ::Rex::Post::PostgreSQL::Ui::Console.new(self)
     super(rstream, opts)
   end
 
@@ -35,7 +31,6 @@ def bootstrap(datastore = {}, handler = nil)
     session = self
     session.init_ui(self.user_input, self.user_output)
 
-    #@info = "SMB #{datastore['USERNAME']} @ #{@peer_info}"
     @info = "PostgreSQL #{datastore['USERNAME']} @ #{@peer_info}"
   end
 
@@ -51,7 +46,6 @@ def type
   # @return [String] The type of the session
   #
   def self.type
-    # 'SMB'
     'PostgreSQL'
   end
 
@@ -78,8 +72,8 @@ def desc
   def init_ui(input, output)
     self.user_input = input
     self.user_output = output
-    console.init_ui(input, output)
-    console.set_log_source(log_source)
+    self.console.init_ui(self.user_input, self.user_output)
+    self.console.set_log_source(self.log_source)
 
     super
   end
@@ -90,8 +84,8 @@ def init_ui(input, output)
   # Resets the console's I/O handles.
   #
   def reset_ui
-    console.unset_log_source
-    console.reset_ui
+    self.console.unset_log_source
+    self.console.reset_ui
   end
 
   ##
@@ -100,10 +94,8 @@ def reset_ui
   # Override the basic session interaction to use shell_read and
   # shell_write instead of operating on rstream directly.
   def _interact
-    framework.events.on_session_interact(self)
-    Rex::Ui::Text::Shell::HistoryManager.with_context(name: self.type.to_sym) {
-      _interact_stream
-    }
+    self.framework.events.on_session_interact(self)
+    self.framework.history_manager.with_context(name: self.type.to_sym) { self._interact_stream }
   end
 
   ##
@@ -112,20 +104,19 @@ def _interact
   def _interact_stream
     self.framework.events.on_session_interact(self)
 
-    console.framework = self.framework
+    self.console.framework = self.framework
     # if framework.datastore['MeterpreterPrompt']
     #   console.update_prompt(framework.datastore['MeterpreterPrompt'])
     # end
     # Call the console interaction of the smb client and
     # pass it a block that returns whether or not we should still be
     # interacting.  This will allow the shell to abort if interaction is
     # canceled.
-    console.interact { self.interacting != true }
-    console.framework = nil
+    self.console.interact { self.interacting != true }
+    self.console.framework = nil
 
     # If the stop flag has been set, then that means the user exited.  Raise
     # the EOFError so we can drop this handle like a bad habit.
-    raise ::EOFError if (console.stopped? == true)
+    raise ::EOFError if (self.console.stopped? == true)
   end
-
-end
+end

lib/rex/post/postgresql.rb

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

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

@@ -4,61 +4,47 @@ module Rex
   module Post
     module PostgreSQL
       module Ui
-
         ###
         #
         # This class provides a shell driven interface to the PostgreSQL client API.
         #
         ###
         class Console
-
           include Rex::Ui::Text::DispatcherShell
 
           # 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/shares'
+          require 'rex/post/postgresql/ui/console/command_dispatcher/db'
 
           #
-          # Initialize the SMB console.
+          # Initialize the PostgreSQL console.
           #
-          # @param [Msf::Sessions::SMB] session
+          # @param [Msf::Sessions::PostgreSQL] session
           def initialize(session)
-            # @type [::Msf::Config.meterpreter_history]
-            todo_hist_file = nil
-            super("%undPostgreSQL%clr", '>', todo_hist_file, nil, :postgresql)
-            #if Rex::Compat.is_windows
-              #super("smb")
-            #  super('postgresql')
-            #else
-            #  todo_hist_file = nil # Msf::Config.meterpreter_history
-              #super("%undSMB%clr", '>', todo_hist_file, nil, :smb)
-            #  super("%undPostgreSQL%clr", '>', todo_hist_file, nil, :postgresql)
-            #end
-
             # The postgresql client context
             self.session = session
             self.client = session.client
+            prompt = "%undPostgreSQL @ #{self.client.conn.remote_address.ip_address}:#{self.client.conn.remote_address.ip_port}%clr"
+            history_manager = self.session.framework&.history_manager&.with_context(name: :postgresql)
+            super(prompt, '>', history_manager, self.session&.framework, :postgresql)
 
             # Queued commands array
             self.commands = []
 
             # Point the input/output handles elsewhere
             reset_ui
 
-            enstack_dispatcher(Rex::Post::SMB::Ui::Console::CommandDispatcher::Core)
-            #enstack_dispatcher(Rex::Post::SMB::Ui::Console::CommandDispatcher::Shares)
+            enstack_dispatcher(::Rex::Post::PostgreSQL::Ui::Console::CommandDispatcher::Core)
+            enstack_dispatcher(::Rex::Post::PostgreSQL::Ui::Console::CommandDispatcher::DB)
 
             # Set up logging to whatever logsink 'core' is using
-            #if ! $dispatcher['smb']
             if ! $dispatcher['postgresql']
-              #$dispatcher['smb'] = $dispatcher['core']
               $dispatcher['postgresql'] = $dispatcher['core']
             end
           end
 
           #
-          # Called when someone wants to interact with the smb client.  It's
           # Called when someone wants to interact with the postgresql client.  It's
           # assumed that init_ui has been called prior.
           #
@@ -103,7 +89,7 @@ def run_command(dispatcher, method, arguments)
               log_error(e.message)
             rescue ::Errno::EPIPE, ::OpenSSL::SSL::SSLError, ::IOError
               self.session.kill
-            rescue ::Exception => e
+            rescue ::StandardError => e
               log_error("Error running command #{method}: #{e.class} #{e}")
               elog(e)
             end

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

@@ -6,15 +6,12 @@ module Rex
   module Post
     module PostgreSQL
       module Ui
-
         ###
         #
-        # Base class for all command dispatchers within the SMB console user
-        # interface.
+        # Base class for all command dispatchers within the PostgreSQL console user interface.
         #
         ###
         module Console::CommandDispatcher
-
           include Rex::Ui::Text::DispatcherShell::CommandDispatcher
 
           #
@@ -50,7 +47,7 @@ def initialize(console)
           end
 
           #
-          # Returns the smb client context.
+          # Returns the PostgreSQL client context.
           #
           # @return [PostgreSQL::Client]
           def client
@@ -59,30 +56,21 @@ def client
           end
 
           #
-          # Returns the smb session context.
+          # Returns the PostgreSQL session context.
           #
           # @return [Msf::Sessions::PostgreSQL]
           def session
             console = shell
             console.session
           end
 
-          # TODO: Rename - required for most of RubySMB's APIs, i.e. "path = "\\\\#{address}\\share"
           def address
-            # TOOD: How do we accurately get the target smb address when pivoting? Also Needs ipv6 support
-            address, port = client.dispatcher.tcp_socket.peerinfo.split(':')
-            address
+            self.client.conn.remote_address.ip_address
           end
 
-          # PostgreSQL doesn't have shares.
-          #
-          # Returns the active share
-          #
-          # @return [RubySMB::SMB2::Tree]
-          #def active_share
-          #  console = shell
-          #  console.active_share
-          #end
+          def port
+            self.client.conn.remote_address.ip_port
+          end
 
           #
           # Returns the commands that meet the requirements
@@ -120,9 +108,9 @@ def docs_dir
           #
           def msf_loaded?
             return @msf_loaded unless @msf_loaded.nil?
-            # if we get here we must not have initialized yet
 
-            @msf_loaded = !!(client.framework)
+            # if we get here we must not have initialized yet
+            @msf_loaded = !!(self.client.framework)
           end
 
           #

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

@@ -86,23 +86,25 @@ def cmd_background_help
           end
 
           def cmd_background
-            # TODO: client is postgresql so doesn't have a name, and session isn't available in this context
-            # TODO: Check this.
-            # print_status("Backgrounding session #{client.name}...")
-            print_status("Backgrounding session...")
-            session.interacting = false
+            print_status("Backgrounding session #{self.session.name}...")
+            self.session.interacting = false
           end
 
           alias cmd_bg cmd_background
           alias cmd_bg_help cmd_background_help
 
           #
-          # Terminates the SMB session.
+          # Terminates the PostgreSQL session.
           #
           def cmd_exit(*args)
-            print_status('Shutting down PostgreSQL... TODO properly disconnect/close socket etc')
-            # client.core.shutdown rescue nil
-            # client.shutdown_passive_dispatcher
+            print_status('Shutting down PostgreSQL...')
+
+            begin
+              self.client.close
+            rescue ::RuntimeError => e # Connection already closed
+              print_warning e.message
+            end
+
             shell.stop
           end
 
@@ -140,10 +142,8 @@ def cmd_irb(*args)
 
             if expressions.empty?
               print_status('Starting IRB shell...')
-              # TODO: What are the semantics of session vs client for SMB?
-              # TODO: Check this.
-              print_status("You are in the \"client\" (session) object\n")
-              Rex::Ui::Text::Shell::HistoryManager.with_context(name: :irb) do
+              print_status("You are in the PostgreSQL command dispatcher object\n")
+              framework.history_manager.with_context(name: :irb) do
                 Rex::Ui::Text::IrbShell.new(session).run
               end
             else