Skip to content

Commit 3f3531d

Browse files
committed
Add test for command shell sessions
1 parent 3b4302d commit 3f3531d

File tree

2 files changed

+177
-15
lines changed

2 files changed

+177
-15
lines changed

lib/msf/base/sessions/command_shell.rb

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ class CommandShell
3030
include Rex::Ui::Text::Resource
3131

3232
@@irb_opts = Rex::Parser::Arguments.new(
33-
'-h' => [false, 'Help menu.' ],
33+
['-h', '--help'] => [false, 'Help menu.' ],
3434
'-e' => [true, 'Expression to evaluate.']
3535
)
3636

@@ -235,35 +235,35 @@ def cmd_sessions_help
235235
end
236236

237237
def cmd_sessions(*args)
238-
if args.length.zero? || args[0].to_i <= 0
239-
# No args
238+
if args.length != 1
239+
print_status "Wrong number of arguments expected: 1, received: #{args.length}"
240240
return cmd_sessions_help
241241
end
242242

243-
if args.length == 1 && (args[1] == '-h' || args[1] == 'help')
244-
# One arg, and args[1] => '-h' '-H' 'help'
243+
if args[0] == '-h' || args[0] == '--help'
245244
return cmd_sessions_help
246245
end
247246

248-
if args.length != 1
249-
# More than one argument
250-
return cmd_sessions_help
247+
session_id = args[0].to_i
248+
if session_id <= 0
249+
print_status 'Invalid session id'
250+
cmd_sessions_help
251251
end
252252

253-
if args[0].to_s == self.name.to_s
253+
if session_id == self.sid
254254
# Src == Dst
255255
print_status("Session #{self.name} is already interactive.")
256256
else
257257
print_status("Backgrounding session #{self.name}...")
258258
# store the next session id so that it can be referenced as soon
259259
# as this session is no longer interacting
260-
self.next_session = args[0]
260+
self.next_session = session_id
261261
self.interacting = false
262262
end
263263
end
264264

265265
def cmd_resource(*args)
266-
if args.empty?
266+
if args.empty? || args[0] == '-h' || args[0] == '--help'
267267
cmd_resource_help
268268
return false
269269
end
@@ -320,9 +320,9 @@ def cmd_shell_help()
320320
end
321321

322322
def cmd_shell(*args)
323-
if args.length == 1 && (args[1] == '-h' || args[1] == 'help')
324-
# One arg, and args[1] => '-h' '-H' 'help'
325-
return cmd_sessions_help
323+
if args.length == 1 && (args[0] == '-h' || args[0] == '--help')
324+
# One arg, and args[0] => '-h' '--help'
325+
return cmd_shell_help
326326
end
327327

328328
if platform == 'windows'
@@ -570,7 +570,7 @@ def cmd_pry_help
570570
# Open the Pry debugger on the current session
571571
#
572572
def cmd_pry(*args)
573-
if args.include?('-h')
573+
if args.include?('-h') || args.include?('--help')
574574
cmd_pry_help
575575
return
576576
end
Lines changed: 162 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,162 @@
1+
# frozen_string_literal: true
2+
3+
require 'spec_helper'
4+
5+
COMMANDS = %i[help background sessions resource shell download upload source irb pry]
6+
7+
RSpec.describe Msf::Sessions::CommandShell do
8+
let(:type) { 'shell' }
9+
10+
it 'should have the correct type' do
11+
expect(described_class.type).to eq(type)
12+
end
13+
14+
it 'should be able to cleanup files' do
15+
expect(described_class.can_cleanup_files).to eq(true)
16+
end
17+
18+
context 'when we have a command shell session' do
19+
subject(:command_shell) { described_class.new(nil) }
20+
let(:command_functions) do
21+
COMMANDS.map { |command| "cmd_#{command}" }
22+
end
23+
let(:command_help_functions) do
24+
command_functions.map { |command| "#{command}_help" }
25+
end
26+
let(:description) { 'Command shell' }
27+
28+
it { is_expected.to respond_to(*command_functions) }
29+
it { is_expected.to respond_to(*command_help_functions) }
30+
31+
it 'should have the correct type' do
32+
expect(subject.type).to eq(type)
33+
end
34+
it 'should have the correct description' do
35+
expect(subject.desc).to eq(description)
36+
end
37+
38+
it 'should not support aborting the process running in the session' do
39+
expect(subject.abort_foreground_supported).to be(true)
40+
end
41+
42+
it 'should initialise the shell by default' do
43+
expect(subject.shell_init).to be(true)
44+
end
45+
46+
describe 'Builtin commands' do
47+
COMMANDS.each do |command|
48+
next if command == :help
49+
50+
describe "#cmd_#{command}" do
51+
context 'when called with the `-h` argument' do
52+
it 'should call the corresponding help function' do
53+
expect(subject).to receive("cmd_#{command}_help")
54+
subject.send("cmd_#{command}", '-h')
55+
end
56+
end
57+
58+
context 'when called with the `--help` argument' do
59+
it 'should call the corresponding help function' do
60+
expect(subject).to receive("cmd_#{command}_help")
61+
subject.send("cmd_#{command}", '--help')
62+
end
63+
end
64+
end
65+
end
66+
end
67+
68+
describe '#run_builtin_cmd' do
69+
COMMANDS.each do |command|
70+
context "when called with `#{command}`" do
71+
it "should call cmd_#{command}" do
72+
expect(subject).to receive("cmd_#{command}")
73+
subject.run_builtin_cmd(command.to_s, nil)
74+
end
75+
end
76+
end
77+
end
78+
79+
describe '#run_single' do
80+
COMMANDS.each do |command|
81+
context "when called with builtin command `#{command}`" do
82+
it 'should call the builtin function' do
83+
expect(subject).to receive(:run_builtin_cmd)
84+
subject.run_single(command.to_s)
85+
end
86+
end
87+
end
88+
89+
context 'when called with a non-builtin command' do
90+
let(:cmd) { 'some_command' }
91+
it 'should write the command to the shell' do
92+
expect(subject).to receive(:shell_write).with("#{cmd}\n")
93+
subject.run_single(cmd)
94+
end
95+
end
96+
end
97+
98+
describe '#process_autoruns' do
99+
let(:initial_auto_run_script) { 'initial_auto_run_script' }
100+
let(:auto_run_script) { 'auto_run_script' }
101+
102+
context 'The datastore is empty' do
103+
let(:datastore) do
104+
Msf::DataStore.new
105+
end
106+
it 'should not execute any script' do
107+
is_expected.not_to receive(:execute_script)
108+
subject.process_autoruns(datastore)
109+
end
110+
end
111+
112+
context 'The datastore contains an `InitialAutoRunScript`' do
113+
let(:datastore) do
114+
datastore = Msf::DataStore.new
115+
datastore['InitialAutoRunScript'] = initial_auto_run_script
116+
datastore
117+
end
118+
it 'should execute the script' do
119+
is_expected.to receive(:execute_script).with(initial_auto_run_script)
120+
subject.process_autoruns(datastore)
121+
end
122+
end
123+
124+
context 'The datastore contains an `AutoRunScript`' do
125+
let(:datastore) do
126+
datastore = Msf::DataStore.new
127+
datastore['AutoRunScript'] = auto_run_script
128+
datastore
129+
end
130+
it 'should execute the script' do
131+
is_expected.to receive(:execute_script).with(auto_run_script)
132+
subject.process_autoruns(datastore)
133+
end
134+
end
135+
136+
context 'The datastore contains both `InitialAutoRunScript` and `AutoRunScript`' do
137+
let(:datastore) do
138+
datastore = Msf::DataStore.new
139+
datastore['InitialAutoRunScript'] = initial_auto_run_script
140+
datastore['AutoRunScript'] = auto_run_script
141+
datastore
142+
end
143+
it 'should execute initial script before the auto run script' do
144+
is_expected.to receive(:execute_script).ordered.with(initial_auto_run_script)
145+
is_expected.to receive(:execute_script).ordered.with(auto_run_script)
146+
subject.process_autoruns(datastore)
147+
end
148+
end
149+
end
150+
151+
context 'when the platform is windows' do
152+
let(:platform) { 'windows' }
153+
before(:each) do
154+
subject.platform = platform
155+
end
156+
157+
it 'should not support aborting the process running in the session' do
158+
expect(subject.abort_foreground_supported).to be(false)
159+
end
160+
end
161+
end
162+
end

0 commit comments

Comments
 (0)