Merge pull request #12 from jvazquez-r7/review_4989

Keep old values when bad file plus specs

lib/msf/core/option_container.rb

@@ -378,17 +378,18 @@ class OptAddressRange < OptBase
 
   def normalize(value)
     return nil unless value.kind_of?(String)
-    if (value =~ /^rand:(.*)/)
+    if value =~ /^rand:(.*)/
       count = $1.to_i
       return false if count < 1
       ret = ''
-      count.times {
-        ret << " " if not ret.empty?
-        ret << [ rand(0x100000000) ].pack("N").unpack("C*").map{|x| x.to_s }.join(".")
-      }
+      count.times do
+        ret << ' ' unless ret.empty?
+        ret << [ rand(0x100000000) ].pack('N').unpack('C*').map{|x| x.to_s }.join('.')
+      end
       return ret
     end
-    return value
+
+    value
   end
 
   def valid?(value)

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

@@ -2068,13 +2068,26 @@ class Core
       return true
     end
 
-    # If the value starts with file: and exists, load the file as the value
-    if value =~ /^file:(.*)/ && ::File.file?($1)
+    # If the value starts with file: exists, and size isn't too big load the file as the value
+    # Otherwise keep the old value
+    if value =~ /^file:(.*)/
       fname = $1
-      if ::File.size(fname) > (1024*1024)
-        print_error("The file name specified is too big (over 1Mb)")
-      else
-        ::File.open(fname, "rb") {|fd| value = fd.read(fd.stat.size) }
+
+      begin
+        fd = ::File.new(fname, 'rb')
+      rescue ::Errno::ENOENT
+        print_error('The file name specified does not exist')
+        value = datastore[name]
+        fd = nil
+      end
+
+      if fd && fd.stat.size > (1024 * 1024)
+        print_error('The file name specified is too big (over 1Mb)')
+        value = datastore[name]
+        fd.close
+      elsif fd
+        value = fd.read(fd.stat.size)
+        fd.close
       end
     end
 

spec/lib/msf/ui/console/command_dispatcher/core_spec.rb

@@ -160,6 +160,70 @@ describe Msf::Ui::Console::CommandDispatcher::Core do
         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
+
+        context "when using file: prefix in the value" do
+          context "when the file exists" do
+
+            before(:each) do
+              allow(::File).to receive(:new) do |filename, mode|
+                fd = StringIO.new(file_contents, mode)
+                fd
+              end
+
+              allow_any_instance_of(::StringIO).to receive(:stat) do |io|
+                file_contents
+              end
+            end
+
+            context "when the size is 1MB" do
+              let(:file_name) do
+                ::Rex::Text.rand_text_alpha(10).upcase
+              end
+
+              let(:file_contents) do
+                ::Rex::Text.rand_text_alpha(1024 * 1024).upcase
+              end
+
+              it "should show the new value" do
+                set_and_test_variable(name, nil, "file:/#{file_name}", nil, /^#{name} => #{file_contents}$/)
+              end
+            end
+
+            context "when the size is greater than 1MB" do
+              let(:file_name) do
+                ::Rex::Text.rand_text_alpha(10).upcase
+              end
+
+              let(:file_contents) do
+                ::Rex::Text.rand_text_alpha(1024 * 1025).upcase
+              end
+
+              it "should show the old value" do
+                set_and_test_variable(name, nil, "file:/#{file_name}", nil, /^#{name} => $/)
+              end
+            end
+
+            context "when the size is less than 1MB" do
+              let(:file_name) do
+                ::Rex::Text.rand_text_alpha(10).upcase
+              end
+
+              let(:file_contents) do
+                ::Rex::Text.rand_text_alpha(10).upcase
+              end
+
+              it "should show the new value" do
+                set_and_test_variable(name, nil, "file:/#{file_name}", nil, /^#{name} => #{file_contents}$/)
+              end
+            end
+          end
+
+          context "when the file doesn't exist" do
+            it "should show the old value" do
+              set_and_test_variable(name, nil, "file:/#{::Rex::Text.rand_text_alpha(10).upcase}", nil, /^#{name} => $/)
+            end
+          end
+        end
       end
     end
   end