Add checks for CSV.load and Marshal.load

and combine with YAML checks

lib/brakeman/checks/check_yaml_load.rb → lib/brakeman/checks/check_deserialize.rb

 require 'brakeman/checks/base_check'
 
-#YAML.load can be used for remote code execution
-class Brakeman::CheckYAMLLoad < Brakeman::BaseCheck
+class Brakeman::CheckDeserialize < Brakeman::BaseCheck
   Brakeman::Checks.add self
 
-  @description = "Checks for uses of YAML.load"
+  @description = "Checks for unsafe deserialization of objects"
 
   def run_check
-    yaml_methods = [:load, :load_documents, :load_stream, :parse_documents, :parse_stream]
+    check_yaml
+    check_csv
+    check_marshal
+  end
+
+  def check_yaml
+    check_methods :YAML, :load, :load_documents, :load_stream, :parse_documents, :parse_stream
+  end
+
+  def check_csv
+    check_methods :CSV, :load
+  end
+
+  def check_marshal
+    check_methods :Marshal, :load, :restore
+  end
 
-    tracker.find_call(:target => :YAML, :methods => yaml_methods ).each do |result|
-      check_yaml_load result
+  def check_methods target, *methods
+    tracker.find_call(:target => target, :methods => methods ).each do |result|
+      check_deserialize result, target
     end
   end
 
-  def check_yaml_load result
+  def check_deserialize result, target, arg = nil
     return if duplicate? result
     add_result result
 
-    arg = result[:call].first_arg
+    arg ||= result[:call].first_arg
     method = result[:call].method
 
     if input = has_immediate_user_input?(arg)
@@ -28,7 +43,7 @@ class Brakeman::CheckYAMLLoad < Brakeman::BaseCheck
     end
 
     if confidence
-      message = "YAML.#{method} called with #{friendly_type_of input}"
+      message = "#{target}.#{method} called with #{friendly_type_of input}"
 
       warn :result => result,
         :warning_type => "Remote Code Execution",
@@ -36,7 +51,7 @@ class Brakeman::CheckYAMLLoad < Brakeman::BaseCheck
         :message => message,
         :user_input => input.match,
         :confidence => confidence,
-        :link_path => "remote_code_execution_yaml_load"
+        :link_path => "unsafe_deserialization"
     end
   end
 end

test/apps/rails3.1/app/controllers/other_controller.rb

@@ -64,4 +64,12 @@ class OtherController < ApplicationController
     sanitize something
     @css = sanitize_css(some_css)
   end
+
+  def test_deserialization
+    CSV.load params[:csv]
+
+    Marshal.load params[:object]
+
+    Marshal.restore User.find(1).cool_stored_thing
+  end
 end

test/tests/test_rails31.rb

@@ -15,7 +15,7 @@ class Rails31Tests < Test::Unit::TestCase
       :model => 3,
       :template => 22,
       :controller => 4,
-      :warning => 68 }
+      :warning => 71 }
   end
 
   def test_without_protection
@@ -969,6 +969,7 @@ class Rails31Tests < Test::Unit::TestCase
       :confidence => 0,
       :file => /admin_controller\.rb/
   end
+
   def test_unsafe_reflection_constantize_indirect
     assert_warning :type => :warning,
       :warning_type => "Remote Code Execution",
@@ -977,4 +978,37 @@ class Rails31Tests < Test::Unit::TestCase
       :confidence => 1,
       :file => /admin_controller\.rb/
   end
+
+  def test_csv_load
+    assert_warning :type => :warning,
+      :warning_code => 25,
+      :fingerprint => "3b58b691bf7ef0b244ee463aa812e4e6ffe3fe1075c8bd138c0cb5a77f365f41",
+      :warning_type => "Remote Code Execution",
+      :line => 69,
+      :message => /^CSV\.load\ called\ with\ parameter\ value/,
+      :confidence => 0,
+      :relative_path => "app/controllers/other_controller.rb"
+  end
+
+  def test_marshal_load
+    assert_warning :type => :warning,
+      :warning_code => 25,
+      :fingerprint => "ecdb984aa40fbe7d42a74ab474a412579b42b36c630bcac640d382e108109437",
+      :warning_type => "Remote Code Execution",
+      :line => 71,
+      :message => /^Marshal\.load\ called\ with\ parameter\ value/,
+      :confidence => 0,
+      :relative_path => "app/controllers/other_controller.rb"
+  end
+
+  def test_marshal_restore
+    assert_warning :type => :warning,
+      :warning_code => 25,
+      :fingerprint => "78ef96a81c8b02f97992a7056e4d9696ab238e12bc8a7a3204df29ef11e0a3fe",
+      :warning_type => "Remote Code Execution",
+      :line => 73,
+      :message => /^Marshal\.restore\ called\ with\ model\ attrib/,
+      :confidence => 1,
+      :relative_path => "app/controllers/other_controller.rb"
+  end
 end