Land rapid7#4504, @disenchant's get/getg improvement

Author: jhart-r7

features/commands/help.feature

@@ -19,6 +19,8 @@ Feature: Help command
           connect       Communicate with a host
           edit          Edit the current module with $VISUAL or $EDITOR
           exit          Exit the console
+          get           Gets the value of a variable
+          getg          Gets the value of a global variable
           go_pro        Launch Metasploit web GUI
           grep          Grep the output of another command
           help          Help menu

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

@@ -15,6 +15,11 @@ def rpc_stop
     self.service.stop
   end
 
+  def rpc_getg(var)
+    val = framework.datastore[var]
+    { var.to_s => val.to_s }
+  end
+
   def rpc_setg(var, val)
     framework.datastore[var] = val
     { "result" => "success" }

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

@@ -115,6 +115,8 @@ def commands
       "color"    => "Toggle color",
       "exit"     => "Exit the console",
       "edit"     => "Edit the current module with $VISUAL or $EDITOR",
+      "get"      => "Gets the value of a variable",
+      "getg"     => "Gets the value of a global variable",
       "go_pro"   => "Launch Metasploit web GUI",
       "grep"     => "Grep the output of another command",
       "help"     => "Help menu",
@@ -2298,6 +2300,81 @@ def cmd_unload_tabs(str, words)
     return tabs
   end
 
+ def cmd_get_help
+    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
+  end
+
+  #
+  # Gets a value if it's been set.
+  #
+  def cmd_get(*args)
+
+    # Figure out if these are global variables
+    global = false
+
+    if (args[0] == '-g')
+      args.shift
+      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 && !global)
+      datastore = active_module.datastore
+    else
+      datastore = framework.datastore
+    end
+
+    args.each { |var| print_line("#{var} => #{datastore[var]}") }
+  end
+
+  #
+  # Tab completion for the get command
+  #
+  # @param str [String] the string currently being typed before tab was hit
+  # @param words [Array<String>] the previously completed words on the command line.  words is always
+  # at least 1 when tab completion has reached this stage since the command itself has been completed
+
+  def cmd_get_tabs(str, words)
+    datastore = active_module ? active_module.datastore : self.framework.datastore
+    datastore.keys
+  end
+
+  def cmd_getg_help
+    print_line "Usage: getg var1 [var2 ...]"
+    print_line
+    print_line "Exactly like get -g, get global variables"
+    print_line
+  end
+
+  #
+  # Gets variables in the global data store.
+  #
+  def cmd_getg(*args)
+    args.unshift('-g')
+
+    cmd_get(*args)
+  end
+
+  #
+  # Tab completion for the getg command
+  #
+  # @param str [String] the string currently being typed before tab was hit
+  # @param words [Array<String>] the previously completed words on the command line.  words is always
+  # at least 1 when tab completion has reached this stage since the command itself has been completed
+
+  def cmd_getg_tabs(str, words)
+    self.framework.datastore.keys
+  end
+
   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