Consolidate session tests

lib/metasploit/framework/login_scanner/postgres.rb

@@ -11,7 +11,7 @@ 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,
+        # @returns [Boolean] If a login is successful and this attribute is true - a Msf::Db::PostgresPR::Connection instance is used as proof,
         #   and the socket is not immediately closed
         attr_accessor :use_client_as_proof
 

lib/msf/base/sessions/mssql.rb

@@ -4,10 +4,6 @@
 
 class Msf::Sessions::MSSQL < Msf::Sessions::Sql
 
-  # @return [String] The address MSSQL is running on
-  attr_accessor :address
-  # @return [Integer] The port MSSQL is running on
-  attr_accessor :port
   attr_reader :framework
 
   def initialize(rstream, opts = {})
@@ -40,18 +36,4 @@ def self.can_cleanup_files
   def desc
     'MSSQL'
   end
-
-  def address
-    return @address if @address
-
-    @address, @port = client.sock.peerinfo.split(':')
-    @address
-  end
-
-  def port
-    return @port if @port
-
-    @address, @port = client.sock.peerinfo.split(':')
-    @port
-  end
 end

lib/msf/base/sessions/mysql.rb

@@ -19,7 +19,7 @@ def bootstrap(datastore = {}, handler = nil)
     session = self
     session.init_ui(user_input, user_output)
 
-    @info = "MySQL #{datastore['USERNAME']} @ #{client.socket.peerinfo}"
+    @info = "MySQL #{datastore['USERNAME']} @ #{client.peerinfo}"
   end
 
   # @return [String] The type of the session
@@ -36,20 +36,4 @@ def self.can_cleanup_files
   def desc
     'MySQL'
   end
-
-  # @return [Object] The peer address
-  def address
-    return @address if @address
-
-    @address, @port = @client.socket.peerinfo.split(':')
-    @address
-  end
-
-  # @return [Object] The peer host
-  def port
-    return @port if @port
-
-    @address, @port = @client.socket.peerinfo.split(':')
-    @port
-  end
 end

lib/msf/base/sessions/postgresql.rb

@@ -6,7 +6,7 @@ class Msf::Sessions::PostgreSQL < Msf::Sessions::Sql
 
   # @param[Rex::IO::Stream] rstream
   # @param [Hash] opts
-  # @param opts [PostgreSQL::Client] :client
+  # @param opts [Msf::Db::PostgresPR::Connection] :client
   def initialize(rstream, opts = {})
     @client = opts.fetch(:client)
     @console = ::Rex::Post::PostgreSQL::Ui::Console.new(self)
@@ -39,18 +39,4 @@ def self.can_cleanup_files
   def desc
     'PostgreSQL'
   end
-
-  def address
-    return @address if @address
-
-    @address, @port = @client.conn.peerinfo.split(':')
-    @address
-  end
-
-  def port
-    return @port if @port
-
-    @address, @port = @client.conn.peerinfo.split(':')
-    @port
-  end
 end

lib/msf/base/sessions/smb.rb

@@ -75,17 +75,11 @@ def desc
   end
 
   def address
-    return @address if @address
-
-    @address, @port = self.client.dispatcher.tcp_socket.peerinfo.split(':')
-    @address
+    @address ||= simple_client.peerhost
   end
 
   def port
-    return @port if @port
-
-    @address, @port = self.client.dispatcher.tcp_socket.peerinfo.split(':')
-    @port
+    @port ||= simple_client.peerport
   end
 
   ##

lib/msf/base/sessions/sql.rb

@@ -73,14 +73,14 @@ def desc
     raise ::NotImplementedError
   end
 
-  # @return [Object] The peer address
+  # @return [String] The peer address
   def address
-    raise ::NotImplementedError
+    client.peerhost
   end
 
-  # @return [Object] The peer host
+  # @return [Integer] The peer port
   def port
-    raise ::NotImplementedError
+    client.peerport
   end
 
   # Initializes the console's I/O handles.

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

@@ -29,7 +29,7 @@ def initialize(session, opts={})
             self.session = session
             self.client = session.client
             envchange = ::Rex::Proto::MSSQL::ClientMixin::ENVCHANGE
-            prompt = "%undMSSQL @ #{client.sock.peerinfo} (#{client.initial_info_for_envchange(envchange: envchange::DATABASE)[:new]})%clr"
+            prompt = "%undMSSQL @ #{client.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/post/mysql/ui/console.rb

@@ -24,8 +24,7 @@ def initialize(session)
             # The mysql client context
             self.session = session
             self.client = session.client
-            self.client.socket ||= self.client.io
-            prompt = "%undMySQL @ #{client.socket.peerinfo} (#{current_database})%clr"
+            prompt = "%undMySQL @ #{client.peerinfo} (#{current_database})%clr"
             history_manager = Msf::Config.mysql_session_history
             super(prompt, '>', history_manager, nil, :mysql)
 

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

@@ -28,7 +28,7 @@ def initialize(session)
             # The postgresql client context
             self.session = session
             self.client = session.client
-            prompt = "%undPostgreSQL @ #{client.conn.peerinfo} (#{current_database})%clr"
+            prompt = "%undPostgreSQL @ #{client.peerinfo} (#{current_database})%clr"
             history_manager = Msf::Config.postgresql_session_history
             super(prompt, '>', history_manager, nil, :postgresql)
 

lib/rex/proto/mssql/client.rb

@@ -68,6 +68,7 @@ def initialize(framework_module, framework, rhost, rport = 1433, proxies = nil)
           @rhost = rhost
           @rport = rport
           @proxies = proxies
+          @current_database = ''
         end
 
         #

lib/rex/proto/smb/simple_client.rb

@@ -19,11 +19,13 @@ class SimpleClient
   DEFAULT_VERSIONS = [1, 2, 3].freeze
 
   # Public accessors
-  attr_accessor :last_error, :server_max_buffer_size, :address, :port
+  attr_accessor :last_error, :server_max_buffer_size
 
   # Private accessors
   attr_accessor :socket, :client, :direct, :shares, :last_share, :versions
 
+  attr_reader :address, :port
+
   # Pass the socket object and a boolean indicating whether the socket is netbios or cifs
   def initialize(socket, direct = false, versions = DEFAULT_VERSIONS, always_encrypt: true, backend: nil, client: nil)
     self.socket = socket
@@ -258,6 +260,19 @@ def negotiated_smb_version
     self.client.negotiated_smb_version || -1
   end
 
+  alias peerhost address
+
+  def peerport
+    port.to_i
+  end
+
+  def peerinfo
+    "#{peerhost}:#{peerport}"
+  end
+
+  private
+
+  attr_writer :address, :port
 end
 end
 end

modules/auxiliary/scanner/mysql/mysql_login.rb

@@ -200,9 +200,7 @@ def int_version(str)
   def session_setup(result, client)
     return unless (result && client)
 
-    rstream = client.socket || client.io
-
-    my_session = Msf::Sessions::MySQL.new(rstream, { client: client })
+    my_session = Msf::Sessions::MySQL.new(client.io, { client: client })
     merging = {
       'USERPASS_FILE' => nil,
       'USER_FILE'     => nil,

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

@@ -4,7 +4,6 @@
 require 'rex/post/mssql'
 
 RSpec.describe Msf::Sessions::MSSQL do
-  let(:rstream) { instance_double(::Rex::Socket) }
   let(:client) { instance_double(Rex::Proto::MSSQL::Client) }
   let(:opts) { { client: client } }
   let(:console_class) { Rex::Post::MSSQL::Ui::Console }
@@ -16,7 +15,7 @@
   let(:description) { 'MSSQL' }
   let(:can_cleanup_files) { false }
   let(:address) { '192.0.2.1' }
-  let(:port) { '1433' }
+  let(:port) { 1433 }
   let(:peer_info) { "#{address}:#{port}" }
   let(:console) do
     console = Rex::Post::MSSQL::Ui::Console.new(session)
@@ -28,127 +27,11 @@
   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(:initial_info_for_envchange).with({ envchange: 1 }).and_return(envchange_result)
-    allow(rstream).to receive(:peerinfo).and_return(peer_info)
+    allow(client).to receive(:peerinfo).and_return(peer_info)
+    allow(client).to receive(:peerport).and_return(port)
+    allow(client).to receive(:peerhost).and_return(address)
   end
 
-  subject(:session) do
-    mssql_session = described_class.new(rstream, opts)
-    mssql_session.user_input = user_input
-    mssql_session.user_output = user_output
-    mssql_session.name = name
-    mssql_session
-  end
-
-  describe '.type' do
-    it 'should have the correct type' do
-      expect(described_class.type).to eq(type)
-    end
-  end
-
-  describe '.can_cleanup_files' do
-    it 'should be able to cleanup files' do
-      expect(described_class.can_cleanup_files).to eq(can_cleanup_files)
-    end
-  end
-
-  describe '#desc' do
-    it 'should have the correct description' do
-      expect(subject.desc).to eq(description)
-    end
-  end
-
-  describe '#type' do
-    it 'should have the correct type' do
-      expect(subject.type).to eq(type)
-    end
-  end
-
-  describe '#initialize' do
-    context 'without a client' do
-      let(:opts) { {} }
-
-      it 'raises a KeyError' do
-        expect { subject }.to raise_exception(KeyError)
-      end
-    end
-    context 'with a client' do
-      it 'does not raise an exception' do
-        expect { subject }.not_to raise_exception
-      end
-    end
-
-    it 'creates a new console' do
-      expect(subject.console).to be_a(console_class)
-    end
-  end
-
-  describe '#bootstrap' do
-    subject { session.bootstrap }
-
-    it 'keeps the sessions user input' do
-      expect { subject }.not_to change(session, :user_input).from(user_input)
-    end
-
-    it 'keeps the sessions user output' do
-      expect { subject }.not_to change(session, :user_output).from(user_output)
-    end
-
-    it 'sets the console input' do
-      expect { subject }.to change(session.console, :input).to(user_input)
-    end
-
-    it 'sets the console output' do
-      expect { subject }.to change(session.console, :output).to(user_output)
-    end
-
-    it 'sets the log source' do
-      expect { subject }.to change(session.console, :log_source).to(log_source)
-    end
-  end
-
-  describe '#reset_ui' do
-    before(:each) do
-      session.bootstrap
-    end
-
-    subject { session.reset_ui }
-
-    it 'keeps the sessions user input' do
-      expect { subject }.not_to change(session, :user_input).from(user_input)
-    end
-
-    it 'keeps the sessions user output' do
-      expect { subject }.not_to change(session, :user_output).from(user_output)
-    end
-
-    it 'resets the console input' do
-      expect { subject }.to change(session.console, :input).from(user_input).to(nil)
-    end
-
-    it 'resets the console output' do
-      expect { subject }.to change(session.console, :output).from(user_output).to(nil)
-    end
-  end
-
-  describe '#exit' do
-    subject { session.exit }
-
-    it 'exits the session' do
-      expect { subject }.to change(session.console, :stopped?).from(false).to(true)
-    end
-  end
-
-  describe '#address' do
-    subject { session.address }
-
-    it { is_expected.to eq(address) }
-  end
-
-  describe '#port' do
-    subject { session.port }
-
-    it { is_expected.to eq(port) }
-  end
+  it_behaves_like 'client session'
 end