Expose MSSQL initial connection info in client

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

@@ -26,7 +26,8 @@ def initialize(session, opts={})
             # The mssql client context
             self.session = session
             self.client = session.client
-            prompt = "%undMSSQL @ #{client.sock.peerinfo} (#{database_name})%clr"
+            envchange = ::Rex::Proto::MSSQL::ClientMixin::ENVCHANGE
+            prompt = "%undMSSQL @ #{client.sock.peerinfo} (#{client.initial_info_for_envchange(envchange: envchange::DATABASE)[:new]})%clr"
             history_manager = Msf::Config.mssql_session_history
             super(prompt, '>', history_manager, nil, :mssql)
 

lib/rex/proto/mssql/client.rb

@@ -42,6 +42,10 @@ class Client
         # @!attribute send_delay
         #   @return [Integer] The delay between sending packets
         attr_accessor :send_delay
+        # @!attribute initial_connection_info
+        #   @return [Hash] Key-value pairs received from the server during the initial MSSQL connection.
+        # See the spec here: https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-tds/b46a581a-39de-4745-b076-ec4dbb7d13ec
+        attr_accessor :initial_connection_info
 
         def initialize(framework_module, framework, rhost, rport = 1433)
           @framework_module       = framework_module
@@ -175,6 +179,7 @@ def mssql_login(user='sa', pass='', db='', domain_name='')
 
             info = {:errors => []}
             info = mssql_parse_reply(resp, info)
+            self.initial_connection_info = info
 
             return false if not info
             return info[:login_ack] ? true : false
@@ -406,6 +411,7 @@ def mssql_login(user='sa', pass='', db='', domain_name='')
 
           info = {:errors => []}
           info = mssql_parse_reply(resp, info)
+          self.initial_connection_info = info
 
           return false if not info
           info[:login_ack] ? true : false
@@ -640,6 +646,17 @@ def powershell_upload_exec(exe, debug=false)
           print_status("Be sure to cleanup #{var_payload}.exe...")
         end
 
+        # @param [ENVCHANGE] envchange The ENVCHANGE type to get the information for.
+        # @return [Hash] Returns a hash of values if the provided type exists.
+        # @return [Hash] Returns the whole connection info if envchange is nil.
+        # @return [Hash] Returns an empty hash if the provided type is not present.
+        def initial_info_for_envchange(envchange: nil)
+          return self.initial_connection_info if envchange.nil?
+          return nil unless (self.initial_connection_info && self.initial_connection_info.is_a?(::Hash))
+
+          self.initial_connection_info[:envs]&.select { |hash| hash[:type] == envchange }&.first || {}
+        end
+
         def address
           rhost
         end

lib/rex/proto/mssql/client_mixin.rb

@@ -32,6 +32,30 @@ module ClientMixin
   STATUS_RESETCONNECTION         = 0x08 # TDS 7.1+
   STATUS_RESETCONNECTIONSKIPTRAN = 0x10 # TDS 7.3+
 
+  # Mappings for ENVCHANGE types
+  # See the TDS Specification here: https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-tds/2b3eb7e5-d43d-4d1b-bf4d-76b9e3afc791
+  module ENVCHANGE
+    DATABASE = 1
+    LANGUAGE = 2
+    CHARACTER_SET = 3
+    PACKET_SIZE = 4
+    UNICODE_LOCAL_ID = 5
+    UNICODE_COMPARISON_FLAGS = 6
+    SQL_COLLATION = 7
+    BEGIN_TRANSACTION = 8
+    COMMIT_TRANSACTION = 9
+    ROLLBACK_TRANSACTION = 10
+    ENLIST_DTC_TRANSACTION = 11
+    DEFECT_TRANSACTION = 12
+    REAL_TIME_LOG_SHIPPING = 13
+    PROMOTE_TRANSACTION = 15
+    TRANSACTION_MANAGER_ADDRESS = 16
+    TRANSACTION_ENDED = 17
+    COMPLETION_ACKNOWLEDGEMENT = 18
+    NAME_OF_USER_INSTANCE = 19
+    ROUTING_INFORMATION = 20
+  end
+
   def mssql_print_reply(info)
     print_status("SQL Query: #{info[:sql]}")
 

spec/lib/msf/base/sessions/mssql_spec.rb

@@ -11,9 +11,6 @@
   let(:user_input) { instance_double(Rex::Ui::Text::Input::Readline) }
   let(:user_output) { instance_double(Rex::Ui::Text::Output::Stdio) }
   let(:name) { 'mssql' }
-  let(:query_result) do
-    { rows: [['mssql']]}
-  end
   let(:log_source) { "session_#{name}" }
   let(:type) { 'mssql' }
   let(:description) { 'MSSQL' }
@@ -26,12 +23,13 @@
     console.disable_output = true
     console
   end
+  let(:envchange_result) { { type: 1, old: 'master', new: 'master' } }
 
   before(:each) do
     allow(user_input).to receive(:intrinsic_shell?).and_return(true)
     allow(user_input).to receive(:output=)
     allow(client).to receive(:sock).and_return(rstream)
-    allow(client).to receive(:mssql_query).with('SELECT DB_NAME();').and_return(query_result)
+    allow(client).to receive(:initial_info_for_envchange).with({ envchange: 1 }).and_return(envchange_result)
     allow(rstream).to receive(:peerinfo).and_return(peer_info)
   end
 

spec/lib/rex/post/mssql/ui/console/command_dispatcher/core_spec.rb

@@ -6,9 +6,6 @@
 RSpec.describe Rex::Post::MSSQL::Ui::Console::CommandDispatcher::Core do
   let(:rstream) { instance_double(::Rex::Socket) }
   let(:client) { instance_double(Rex::Proto::MSSQL::Client) }
-  let(:query_result) do
-    { rows: [['mssql']]}
-  end
   let(:session) { Msf::Sessions::MSSQL.new(nil, { client: client }) }
   let(:address) { '192.0.2.1' }
   let(:port) { '1433' }
@@ -18,10 +15,11 @@
     console.disable_output = true
     console
   end
+  let(:envchange_result) { { type: 1, old: 'master', new: 'master' } }
 
   before(:each) do
     allow(client).to receive(:sock).and_return(rstream)
-    allow(client).to receive(:mssql_query).with('SELECT DB_NAME();').and_return(query_result)
+    allow(client).to receive(:initial_info_for_envchange).with({ envchange: 1 }).and_return(envchange_result)
     allow(rstream).to receive(:peerinfo).and_return(peer_info)
     allow(session).to receive(:client).and_return(client)
     allow(session).to receive(:console).and_return(console)