Fix all the tests

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

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
 

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

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)

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

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

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