Mass Assignment Security Generator

mass_assignment_security_generator.rb

# Copyright (C) 2012 Jason Nochlin
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
# CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
class MassAssignmentSecurityGenerator < Rails::Generators::Base
  source_root File.expand_path('../templates', __FILE__)
  argument :list_type, :type => :string, :default => 'white'
  class_option :include_foreign_keys, :type => :boolean, :default => false, :description => "Include foreign keys in attributes list"

  
  desc "Adds attr_accessible/attr_protected to all models for all their columns"
  def inject_models_with_mass_assignment_security
    directive = case list_type
    when 'white'
      'attr_accessible'
    when 'black'
      'attr_protected'
    end

    files = Dir[Rails.root+'app/models/**/*.rb']

    files.each { |file| require(file) rescue nil }

    ActiveRecord::Base.subclasses.each do |klass|
      underscored = klass.to_s.underscore
      file = files.select { |i| i.match("/#{underscored}.rb") }.first

      if file.nil?
        puts "Couldn't find file for: #{klass}"
        next
      end

      contents = File.read(file)

      if contents.match(/attr_accessible|attr_protected/)
        next
      else
        code = "\n  # Mass Assignment Security"

        columns = klass.columns.map(&:name).reverse.select { |i| options.include_foreign_keys? || !i.match(/_id$/) }
        while (chunk = columns.pop(4).reverse).any?
          code << "\n  #{directive} #{chunk.map { |i| ':' + i }.join(', ')}"
        end
        code << "\n"
        inject_into_class(file, klass.to_s, code)
      end
    end
  end
end