From b40fd2a868d010e52ffffd9b030a13597cf393fb Mon Sep 17 00:00:00 2001
From: Matt Bessey <mbessey@caring.com>
Date: Wed, 25 May 2016 10:52:43 -0700
Subject: [PATCH] [strategies] support optional 2nd instance argument (#40)

As discussed in https://github.com/IFTTT/polo/issues/39 its nice to have the option to access the context (AR instance) from the scrub strategy. this allows that.
---
 README.md               | 16 +++++++++++-----
 lib/polo/translator.rb  |  9 +++++----
 spec/translator_spec.rb | 10 ++++++++++
 3 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index 39b5e7e..1770adb 100644
--- a/README.md
+++ b/README.md
@@ -137,7 +137,7 @@ Warning: This is not a security feature. Fields can still easily be rearranged b
 
 #### Advanced Obfuscation
 
-For more advanced obfuscation, you can pass in a custom obfuscation strategy. Polo will take in a lambda that can be used to transform sensitive data. 
+For more advanced obfuscation, you can pass in a custom obfuscation strategy. Polo will take in a lambda that can be used to transform sensitive data.
 
 Using a `:symbol` as an obfuscate key targets all columns of that name. Passing an SQL selector as a `String` will target columns within the specified table.
 
@@ -148,14 +148,21 @@ Polo.configure do
     first_part = email.split("@")[0]
     "#{first_part}@test.com"
   end
-  
+
   credit_card_strategy = lambda do |credit_card|
     "4123 4567 8910 1112"
   end
-  
+
+  # If you need the context of the record for its fields, it is accessible
+  # in the second argument of the strategy
+  social_security_strategy = lambda do |ssn, instance|
+    sprintf("%09d", instance.id)
+  end
+
   obfuscate({
     'chefs.email' => email_strategy, # This only applies to the "email" column in the "chefs" table
-    :credit_card  => credit_card_strategy # This applies to any column named "credit_card" across every table
+    :credit_card  => credit_card_strategy, # This applies to any column named "credit_card" across every table
+    :ssn_strategy => social_security_strategy
   })
 end
 
@@ -200,4 +207,3 @@ $ bundle exec appraisal rake
 ## License
 
 The gem is available as open source under the terms of the [MIT License](http://opensource.org/licenses/MIT).
-
diff --git a/lib/polo/translator.rb b/lib/polo/translator.rb
index 66af832..3cbaa45 100644
--- a/lib/polo/translator.rb
+++ b/lib/polo/translator.rb
@@ -46,8 +46,8 @@ def obfuscate!(instances, fields)
 
           correct_table = table.nil? || instance.class.table_name == table
 
-          if correct_table && value = instance.attributes[field]
-            instance.send("#{field}=", new_field_value(field, strategy, value))
+          if correct_table && instance.attributes[field]
+            instance.send("#{field}=", new_field_value(field, strategy, instance))
           end
         end
       end
@@ -65,11 +65,12 @@ def intersection(attrs, fields)
       attrs & fields.map { |pair| field_name(pair.first) }
     end
 
-    def new_field_value(field, strategy, value)
+    def new_field_value(field, strategy, instance)
+      value = instance.attributes[field]
       if strategy.nil?
         value.split("").shuffle.join
       else
-        strategy.call(value)
+        strategy.arity == 1 ? strategy.call(value) : strategy.call(value, instance)
       end
     end
   end
diff --git a/spec/translator_spec.rb b/spec/translator_spec.rb
index e6dfb17..9b12508 100644
--- a/spec/translator_spec.rb
+++ b/spec/translator_spec.rb
@@ -47,6 +47,16 @@
         end
       end
 
+      context "custom obfuscation strategy using instance context" do
+        let(:obfuscated_fields) do
+          { email: lambda { |field, instance| "#{instance.name}@example.com" } }
+        end
+
+        it "replaces contents of field according to the supplied lambda" do
+          expect(netto.email.to_s).to eq "Netto@example.com"
+        end
+      end
+
       context "no strategy passed in" do
         let(:obfuscated_fields) { [:email] }