From 6762a7b147e2883ba482bbd27750243a305aa6f6 Mon Sep 17 00:00:00 2001 From: dwelch-r7 Date: Wed, 11 Mar 2020 17:06:37 +0000 Subject: [PATCH] Fix all the tests --- lib/msf/base/simple/auxiliary.rb | 20 ++++---- lib/msf/base/simple/exploit.rb | 14 +++--- lib/msf/base/simple/framework.rb | 2 +- ...status_tracker.rb => noop_job_listener.rb} | 6 +-- .../core/rpc/v10/rpc_job_status_tracker.rb | 49 ++++++++++--------- lib/msf/core/rpc/v10/rpc_module.rb | 10 ++-- .../rpc/v10/rpc_job_status_tracker_spec.rb | 48 +++++++----------- 7 files changed, 69 insertions(+), 80 deletions(-) rename lib/msf/base/simple/{noop_job_status_tracker.rb => noop_job_listener.rb} (76%) diff --git a/lib/msf/base/simple/auxiliary.rb b/lib/msf/base/simple/auxiliary.rb index 5b01bf326f..38afcc5bc5 100644 --- a/lib/msf/base/simple/auxiliary.rb +++ b/lib/msf/base/simple/auxiliary.rb @@ -41,7 +41,7 @@ module Auxiliary # Whether or not the exploit should be run in the context of a background # job. # - def self.run_simple(omod, opts = {}, job_status_tracker: NoopJobStatusTracker.instance, &block) + def self.run_simple(omod, opts = {}, job_listener: NoopJobListener.instance, &block) # Clone the module to prevent changes to the original instance mod = omod.replicant @@ -69,8 +69,8 @@ module Auxiliary end run_uuid = Rex::Text.rand_text_alphanumeric(24) - job_status_tracker.waiting run_uuid - ctx = [mod, run_uuid, job_status_tracker] + job_listener.waiting run_uuid + ctx = [mod, run_uuid, job_listener] if(mod.passive? or opts['RunAsJob']) mod.job_id = mod.framework.jobs.start_bg_job( "Auxiliary: #{mod.refname}", @@ -108,7 +108,7 @@ module Auxiliary # # The local output through which data can be displayed. # - def self.check_simple(mod, opts, job_status_tracker: NoopJobStatusTracker.instance) + def self.check_simple(mod, opts, job_listener: NoopJobListener.instance) Msf::Simple::Framework.simplify_module(mod, false) mod._import_extra_options(opts) @@ -122,8 +122,8 @@ module Auxiliary run_uuid = Rex::Text.rand_text_alphanumeric(24) - job_status_tracker.waiting run_uuid - ctx = [mod, run_uuid, job_status_tracker] + job_listener.waiting run_uuid + ctx = [mod, run_uuid, job_listener] if opts['RunAsJob'] mod.job_id = mod.framework.jobs.start_bg_job( @@ -165,16 +165,16 @@ protected def self.job_run_proc(ctx, &block) mod = ctx[0] run_uuid = ctx[1] - job_status_tracker = ctx[2] + job_listener = ctx[2] begin begin - job_status_tracker.start run_uuid + job_listener.start run_uuid mod.setup mod.framework.events.on_module_run(mod) result = block.call(mod) - job_status_tracker.completed(run_uuid, result, mod) + job_listener.completed(run_uuid, result, mod) rescue ::Exception => e - job_status_tracker.failed(run_uuid, e, mod) + job_listener.failed(run_uuid, e, mod) raise end rescue Msf::Auxiliary::Complete diff --git a/lib/msf/base/simple/exploit.rb b/lib/msf/base/simple/exploit.rb index d660f1fb06..11b90cc6c1 100644 --- a/lib/msf/base/simple/exploit.rb +++ b/lib/msf/base/simple/exploit.rb @@ -177,7 +177,7 @@ module Exploit # # The local output through which data can be displayed. # - def self.check_simple(mod, opts, job_status_tracker: NoopJobStatusTracker.instance) + def self.check_simple(mod, opts, job_listener: NoopJobListener.instance) Msf::Simple::Framework.simplify_module(mod, false) mod._import_extra_options(opts) @@ -190,8 +190,8 @@ module Exploit mod.validate run_uuid = Rex::Text.rand_text_alphanumeric(24) - job_status_tracker.waiting run_uuid - ctx = [mod, run_uuid, job_status_tracker] + job_listener.waiting run_uuid + ctx = [mod, run_uuid, job_listener] if opts['RunAsJob'] mod.job_id = mod.framework.jobs.start_bg_job( @@ -218,14 +218,14 @@ module Exploit def self.job_check_proc(ctx) mod = ctx[0] run_uuid = ctx[1] - job_status_tracker = ctx[2] + job_listener = ctx[2] begin - job_status_tracker.start run_uuid + job_listener.start run_uuid mod.setup result = mod.has_check? ? mod.check : Msf::Exploit::CheckCode::Unsupported - job_status_tracker.completed(run_uuid, result, mod) + job_listener.completed(run_uuid, result, mod) rescue => e - job_status_tracker.failed(run_uuid, e, mod) + job_listener.failed(run_uuid, e, mod) mod.handle_exception e end diff --git a/lib/msf/base/simple/framework.rb b/lib/msf/base/simple/framework.rb index 70d4733169..498ae2c3fe 100644 --- a/lib/msf/base/simple/framework.rb +++ b/lib/msf/base/simple/framework.rb @@ -1,7 +1,7 @@ # -*- coding: binary -*- require 'msf/base/simple' require 'msf/base/simple/framework/module_paths' -require 'msf/base/simple/noop_job_status_tracker' +require 'msf/base/simple/noop_job_listener' module Msf module Simple diff --git a/lib/msf/base/simple/noop_job_status_tracker.rb b/lib/msf/base/simple/noop_job_listener.rb similarity index 76% rename from lib/msf/base/simple/noop_job_status_tracker.rb rename to lib/msf/base/simple/noop_job_listener.rb index a3c0a76a61..07a3e1d972 100644 --- a/lib/msf/base/simple/noop_job_status_tracker.rb +++ b/lib/msf/base/simple/noop_job_listener.rb @@ -1,6 +1,6 @@ require 'singleton' -class NoopJobStatusTracker +class NoopJobListener include Singleton @@ -10,10 +10,10 @@ class NoopJobStatusTracker def start(id) end - def completed(id, result, ttl=nil) + def completed(id, result, mod) end - def failed(id, error, ttl=nil) + def failed(id, error, mod) end def running?(id) diff --git a/lib/msf/core/rpc/v10/rpc_job_status_tracker.rb b/lib/msf/core/rpc/v10/rpc_job_status_tracker.rb index 17c2580ea9..0182400511 100644 --- a/lib/msf/core/rpc/v10/rpc_job_status_tracker.rb +++ b/lib/msf/core/rpc/v10/rpc_job_status_tracker.rb @@ -4,11 +4,11 @@ class RpcJobStatusTracker include MonitorMixin - def initialize(result_ttl=nil) + def initialize @ready = Set.new @running = Set.new # Can be expanded upon later to allow the option of a MemCacheStore being backed by redis for example - @results = ResultsMemoryStore.new(expires_in: result_ttl || 5.minutes) + @results = ResultsMemoryStore.new end def waiting(id) @@ -20,12 +20,12 @@ class RpcJobStatusTracker ready.delete(id) end - def completed(id, result, mod, ttl=nil) - add_result(id, {result: result}, mod, ttl) + def completed(id, result, mod) + add_result(id, { result: result }, mod) end - def failed(id, error, mod, ttl=nil) - add_result( id,{error: error.to_s}, mod, ttl) + def failed(id, error, mod) + add_result(id, { error: error.to_s }, mod) end def running?(id) @@ -67,28 +67,33 @@ class RpcJobStatusTracker results.data end - alias :ack :delete + alias ack delete private - def add_result(id, result, mod, ttl=nil) - begin - # ttl of nil means it will take the default expiry time - string = result.to_json - results.write(id, string, ttl) - rescue Exception => e - wlog("Job with id: #{id} finished but the result could not be stored") - wlog("#{e.class}, #{e.message}") - add_fallback_result(id, mod, ttl) - ensure - running.delete(id) - end + def add_result(id, result, mod) + string = result.to_json + results.write(id, string) + rescue ::Exception => e + wlog("Job with id: #{id} finished but the result could not be stored") + wlog("#{e.class}, #{e.message}") + add_fallback_result(id, mod) + ensure + running.delete(id) + end - def add_fallback_result(id, mod, ttl) + def add_fallback_result(id, mod) begin - results.write(id, {error: { message: 'Job finished but the result could not be stored'}, data: { mod: mod.fullname }}, ttl) - rescue Exception => e + string = { + error: { + message: 'Job finished but the result could not be stored' + }, + data: { mod: mod.fullname } + }.to_json + puts string + results.write(id, string) + rescue ::Exception => e wlog("Job with id: #{id} fallback result failed to be stored") wlog("#{e.class}, #{e.message}") end diff --git a/lib/msf/core/rpc/v10/rpc_module.rb b/lib/msf/core/rpc/v10/rpc_module.rb index 2952e75a26..b02d7ca953 100644 --- a/lib/msf/core/rpc/v10/rpc_module.rb +++ b/lib/msf/core/rpc/v10/rpc_module.rb @@ -485,10 +485,6 @@ class RPC_Module < RPC_Base def rpc_execute(mtype, mname, opts) mod = _find_module(mtype,mname) - require 'pry'; binding.pry - - # listener = self.settings.listener - case mtype when 'exploit' _run_exploit(mod, opts) @@ -749,7 +745,7 @@ private 'Action' => opts['ACTION'], 'RunAsJob' => true, 'Options' => opts - }, job_status_tracker: self.job_status_tracker) + }, job_listener: self.job_status_tracker) { "job_id" => job, "uuid" => uuid @@ -760,7 +756,7 @@ private uuid, job = Msf::Simple::Exploit.check_simple(mod,{ 'RunAsJob' => true, 'Options' => opts - }, job_status_tracker: self.job_status_tracker) + }, job_listener: self.job_status_tracker) { "job_id" => job, "uuid" => uuid @@ -772,7 +768,7 @@ private 'Action' => opts['ACTION'], 'RunAsJob' => true, 'Options' => opts - }, job_status_tracker: self.job_status_tracker) + }, job_listener: self.job_status_tracker) { "job_id" => job, "uuid" => uuid diff --git a/spec/lib/msf/core/rpc/v10/rpc_job_status_tracker_spec.rb b/spec/lib/msf/core/rpc/v10/rpc_job_status_tracker_spec.rb index 2cebd14ca9..265d8c4062 100644 --- a/spec/lib/msf/core/rpc/v10/rpc_job_status_tracker_spec.rb +++ b/spec/lib/msf/core/rpc/v10/rpc_job_status_tracker_spec.rb @@ -4,34 +4,13 @@ require 'msf/core/rpc/v10/rpc_job_status_tracker' RSpec.describe RpcJobStatusTracker do - context "With a 1 second ttl " do - let(:job_status_tracker) { described_class.new(result_ttl=1) } - let(:job_id) { "super_random_job_id" } - - - before(:each) do - allow(ActiveSupport::Cache::MemoryStore).to receive(:new).and_call_original - end - - context "A job has completed" do - before(:each) do - generic_result = "This is a job that has finished" - job_status_tracker.completed(job_id, generic_result) - end - - it 'should show as finished' do - expect(job_status_tracker).to be_finished(job_id) - end - - it { expect(ActiveSupport::Cache::MemoryStore).to have_received(:new).with(expires_in: 1) } - end - end - context "With default options" do let(:job_status_tracker) { described_class.new } let(:job_id) { "super_random_job_id" } let(:good_result) { 'yay_success' } let(:bad_result) { 'boo_fail' } + let(:mod) { double("mod") } + context 'A job is waiting' do before(:each) do @@ -75,7 +54,7 @@ RSpec.describe RpcJobStatusTracker do context "The job completes successfully" do before(:each) do - job_status_tracker.completed(job_id, good_result) + job_status_tracker.completed(job_id, good_result, mod) end it 'should not show as waiting' do @@ -94,7 +73,7 @@ RSpec.describe RpcJobStatusTracker do end it 'should have a retrievable result' do - expect(job_status_tracker.result job_id).to eq({result: good_result}) + expect(job_status_tracker.result job_id).to eql({"result" => good_result}) end context "The job is acknowledged" do @@ -121,7 +100,7 @@ RSpec.describe RpcJobStatusTracker do context "The job fails" do before(:each) do - job_status_tracker.failed(job_id, bad_result) + job_status_tracker.failed(job_id, bad_result, mod) end it 'should not show as waiting' do @@ -140,7 +119,7 @@ RSpec.describe RpcJobStatusTracker do end it 'should have a retrievable result' do - expect(job_status_tracker.result job_id).to eq({error: bad_result}) + expect(job_status_tracker.result job_id).to eql({"error" => bad_result}) end context "The job is acknowledged" do @@ -166,8 +145,14 @@ RSpec.describe RpcJobStatusTracker do end context 'The job result is not serializable' do + let(:mock_result) { {mock: "result"} } before(:each) do - job_status_tracker.completed(job_id, proc {"procs can't be serialized"}) + allow(mod).to receive(:fullname).and_return("module_name") + + allow(job_status_tracker.instance_variable_get(:@results)).to receive(:write).with(job_id, mock_result.to_json).and_raise Exception, 'Intentional explosion' + allow(job_status_tracker.instance_variable_get(:@results)).to receive(:write).with(job_id, /error/).and_call_original + + job_status_tracker.completed(job_id, mock_result, mod) end it 'should show as finished' do @@ -175,8 +160,11 @@ RSpec.describe RpcJobStatusTracker do expect(job_status_tracker.result_ids).to eql([job_id]) end - it 'should have an unexpected_error result' do - expect(job_status_tracker.result job_id).to have_key(:unexpected_error) + it 'should have an error result' do + puts job_status_tracker.result job_id + expect(job_status_tracker.result job_id).to have_key("error") + expect(job_status_tracker.result job_id).to have_key("data") + end end end