Add specs for interactive_shell_read and meterpreter_read` methods,…

… refactored to remove code duplication

lib/msf/core/rpc/v10/rpc_session.rb

@@ -209,12 +209,7 @@ def rpc_interactive_shell_read(sid)
     # and we don't know exact type of the session, so we just pass empty string
     session = _valid_session(sid, '')
 
-    unless session.user_output.respond_to?(:dump_buffer)
-      session.init_ui(Rex::Ui::Text::Input::Buffer.new, Rex::Ui::Text::Output::Buffer.new)
-    end
-
-    data = session.user_output.dump_buffer
-    { 'data' => data }
+    _interactive_read(session)
   end
 
   # Reads the output from a meterpreter session (such as a command output).
@@ -233,12 +228,7 @@ def rpc_interactive_shell_read(sid)
   def rpc_meterpreter_read( sid)
     s = _valid_session(sid,"meterpreter")
 
-    if not s.user_output.respond_to? :dump_buffer
-      s.init_ui(Rex::Ui::Text::Input::Buffer.new, Rex::Ui::Text::Output::Buffer.new)
-    end
-
-    data = s.user_output.dump_buffer
-    { "data" => data }
+    _interactive_read(session)
   end
 
 
@@ -560,6 +550,15 @@ def _valid_session(sid,type)
     s
   end
 
+  def _interactive_read(session)
+    unless session.user_output.respond_to?(:dump_buffer)
+      session.init_ui(Rex::Ui::Text::Input::Buffer.new, Rex::Ui::Text::Output::Buffer.new)
+    end
+
+    data = session.user_output.dump_buffer
+    { 'data' => data }
+  end
+
 end
 end
 end

spec/lib/msf/core/rpc/v10/rpc_session_spec.rb

@@ -0,0 +1,121 @@
+# -*- coding:binary -*-
+require 'spec_helper'
+
+RSpec.describe Msf::RPC::RPC_Session do
+  include_context 'Msf::Simple::Framework'
+  include_context 'Metasploit::Framework::Spec::Constants cleaner'
+  include_context 'Msf::Framework#threads cleaner', verify_cleanup_required: false
+  include_context 'wait_for_expect'
+
+  let(:service) { Msf::RPC::Service.new(framework) }
+  let(:rpc) { described_class.new(service) }
+  let(:target_sid) { 1 }
+
+  let(:sessions) do
+    { target_sid => session }
+  end
+
+  before do
+    allow(framework).to receive(:sessions).and_return(sessions)
+  end
+
+  let(:user_output) do
+    output = instance_double(Rex::Ui::Text::Output::Buffer)
+    allow(output).to receive(:dump_buffer).and_return(expected_data)
+    output
+  end
+
+  context 'with meterpreter session' do
+    let(:session) do
+      instance_double(
+        ::Msf::Sessions::Meterpreter_x64_Win,
+        type: 'meterpreter',
+        sid: target_sid
+      )
+    end
+
+    describe '#rpc_meterpreter_read' do
+      let(:expected_data) { 'test response' }
+
+      shared_examples 'meterpreter read' do
+        subject(:response) do
+          rpc.rpc_meterpreter_read(target_sid)
+        end
+
+        it 'returns expected data' do
+          expect(response).to eq({ 'data' => expected_data })
+        end
+      end
+
+      let(:init_ui_handler) do
+        proc do
+          allow(session).to receive(:user_output).and_return(user_output)
+        end
+      end
+
+      context 'when UI is not initialized' do
+        before do
+          allow(session).to receive(:user_output).once.and_return(nil)
+          allow(session).to receive(:init_ui).and_invoke(init_ui_handler)
+        end
+
+        it_behaves_like 'meterpreter read'
+      end
+
+      context 'when UI is initialized' do
+        before do
+          allow(session).to receive(:user_output).and_return(user_output)
+        end
+
+        it_behaves_like 'meterpreter read'
+      end
+    end
+  end
+
+  context 'with postgres session' do
+    let(:session) do
+      instance_double(
+        ::Msf::Sessions::PostgreSQL,
+        type: 'postgres',
+        sid: target_sid
+      )
+    end
+
+    describe '#rpc_interactive_read' do
+      let(:expected_data) { 'test response' }
+
+      shared_examples 'interactive read' do
+        subject(:response) do
+          rpc.rpc_interactive_shell_read(target_sid)
+        end
+
+        it 'returns expected data' do
+          expect(response).to eq({ 'data' => expected_data })
+        end
+      end
+
+      let(:init_ui_handler) do
+        proc do
+          allow(session).to receive(:user_output).and_return(user_output)
+        end
+      end
+
+      context 'when UI is not initialized' do
+        before do
+          allow(session).to receive(:user_output).once.and_return(nil)
+          allow(session).to receive(:init_ui).and_invoke(init_ui_handler)
+        end
+
+        it_behaves_like 'interactive read'
+      end
+
+      context 'when UI is initialized' do
+        before do
+          allow(session).to receive(:user_output).and_return(user_output)
+        end
+
+        it_behaves_like 'interactive read'
+      end
+    end
+  end
+end