Merge pull request #1 from jhart-r7/landing-4504-jhart

Unit tests for rapid7#4504

Author: disenchant

lib/msf/ui/console/command_dispatcher/core.rb

@@ -2298,7 +2298,7 @@ def cmd_unload_tabs(str, words)
   end
 
  def cmd_get_help
-    print_line "Usage: get var1 var2 var3"
+    print_line "Usage: get var1 [var2 ...]"
     print_line
     print_line "The get command is used to get the value of one or more variables."
     print_line
@@ -2317,22 +2317,20 @@ def cmd_get(*args)
       global = true
     end
 
+    # No arguments?  No cookie.
+    if args.empty?
+      global ? cmd_getg_help : cmd_get_help
+      return false
+    end
+
     # Determine which data store we're operating on
-    if (active_module and global == false)
+    if (active_module && !global)
       datastore = active_module.datastore
     else
       datastore = framework.datastore
     end
 
-    # No arguments?  No cookie.
-    if (args.length == 0)
-      cmd_get_help
-      return false
-    end
-
-    while ((val = args.shift))
-      print_line("#{val} => #{datastore[val]}")
-    end
+    args.each { |var| print_line("#{var} => #{datastore[var]}") }
   end
 
   #
@@ -2347,7 +2345,6 @@ def cmd_get_tabs(str, words)
     datastore.keys
   end
 
-
   def cmd_getg_help
     print_line "Usage: getg var1 [var2 ...]"
     print_line
@@ -2375,8 +2372,6 @@ def cmd_getg_tabs(str, words)
     self.framework.datastore.keys
   end
 
-  alias cmd_getg_help cmd_get_help
-
   def cmd_unset_help
     print_line "Usage: unset [-g] var1 var2 var3 ..."
     print_line

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

@@ -0,0 +1,28 @@
+# -*- coding:binary -*-
+require 'spec_helper'
+
+require 'msf/core/rpc/v10/rpc_base'
+require 'msf/core/rpc/v10/rpc_core'
+require 'msf/core/rpc/v10/service'
+
+describe Msf::RPC::RPC_Core do
+  include_context 'Msf::Simple::Framework'
+
+  let(:service) do
+    Msf::RPC::Service.new(framework)
+  end
+
+  let(:core) do
+    Msf::RPC::RPC_Core.new(service)
+  end
+
+  describe '#rpc_getg' do
+    it 'should show an empty value if the variable is unset' do
+      expect(core.rpc_getg('FOO')).to eq({'FOO' => ''})
+    end
+    it 'should show the correct value if the variable is set' do
+      core.rpc_setg('FOO', 'BAR')
+      expect(core.rpc_getg('FOO')).to eq({'FOO' => 'BAR'})
+    end
+  end
+end

spec/lib/msf/ui/command_dispatcher/auxiliary_spec.rb renamed to spec/lib/msf/ui/console/command_dispatcher/auxiliary_spec.rb

@@ -95,4 +95,72 @@ def cell(table, row, column)
       end
     end
   end
+
+  it { is_expected.to respond_to :cmd_get }
+  it { is_expected.to respond_to :cmd_getg }
+
+  def set_and_test_variable(name, framework_value, module_value, framework_re, module_re)
+    # set the current module
+    allow(core).to receive(:active_module).and_return(mod)
+    # always assume set variables validate (largely irrelevant because ours are random)
+    allow(driver).to receive(:on_variable_set).and_return(true)
+    # the specified global value
+    core.cmd_setg(name, framework_value) if framework_value
+    # set the specified local value
+    core.cmd_set(name, module_value) if module_value
+
+    # test the global value if specified
+    if framework_re
+      @output = []
+      core.cmd_getg(name)
+      @output.join.should =~ framework_re
+    end
+
+    # test the local value if specified
+    if module_re
+      @output = []
+      core.cmd_get(name)
+      @output.join.should =~ module_re
+    end
+  end
+
+  describe "#cmd_get and #cmd_getg" do
+    describe "without arguments" do
+      it "should show the correct help message" do
+        core.cmd_get
+        @output.join.should =~ /Usage: get /
+        @output = []
+        core.cmd_getg
+        @output.join.should =~ /Usage: getg /
+      end
+    end
+
+    describe "with arguments" do
+      let(:name) { ::Rex::Text.rand_text_alpha(10).upcase }
+
+      context "with an active module" do
+        let(:mod) do
+          mod = ::Msf::Module.new
+          mod.send(:initialize, {})
+          mod
+        end
+
+        it "should show no value if not set in the framework or module" do
+          set_and_test_variable(name, nil, nil, /^#{name} => $/, /^#{name} => $/)
+        end
+
+        it "should show the correct value when only the module has this variable" do
+          set_and_test_variable(name, nil, 'MODULE', /^#{name} => $/, /^#{name} => MODULE$/)
+        end
+
+        it "should show the correct value when only the framework has this variable" do
+          set_and_test_variable(name, 'FRAMEWORK', nil, /^#{name} => FRAMEWORK$/, /^#{name} => $/)
+        end
+
+        it "should show the correct value when both the module and the framework have this variable" do
+          set_and_test_variable(name, 'FRAMEWORK', 'MODULE', /^#{name} => FRAMEWORK$/, /^#{name} => MODULE$/)
+        end
+      end
+    end
+  end
 end