dlp: Add sample for reid w/ fpe using surrogate type and unwrapped security key (#4051)

* add code sample and test for reid w/ fpe using surrogate type and unwrapped security key

* refactor reidentify_config

Author: lxhfirenking

dlp/deid.py

@@ -394,6 +394,86 @@ def reidentify_with_fpe(
 
 # [END dlp_reidentify_fpe]
 
+# [START dlp_reidentify_free_text_with_fpe_using_surrogate]
+def reidentify_free_text_with_fpe_using_surrogate(
+    project,
+    input_str,
+    alphabet="NUMERIC",
+    surrogate_type="PHONE_TOKEN",
+    unwrapped_key="YWJjZGVmZ2hpamtsbW5vcA==",
+):
+    """Uses the Data Loss Prevention API to reidentify sensitive data in a
+    string that was encrypted by Format Preserving Encryption (FPE) with
+    surrogate type. The encryption is performed with an unwrapped key.
+    Args:
+        project: The Google Cloud project id to use as a parent resource.
+        input_str: The string to deidentify (will be treated as text).
+        alphabet: The set of characters to replace sensitive ones with. For
+            more information, see https://cloud.google.com/dlp/docs/reference/
+            rest/v2beta2/organizations.deidentifyTemplates#ffxcommonnativealphabet
+        surrogate_type: The name of the surrogate custom info type to used
+            during the encryption process.
+        unwrapped_key: The base64-encoded AES-256 key to use.
+    Returns:
+        None; the response from the API is printed to the terminal.
+    """
+    # Import the client library
+    import google.cloud.dlp
+
+    # Instantiate a client
+    dlp = google.cloud.dlp_v2.DlpServiceClient()
+
+    # Convert the project id into a full resource id.
+    parent = dlp.project_path(project)
+
+    # The unwrapped key is base64-encoded, but the library expects a binary
+    # string, so decode it here.
+    import base64
+
+    unwrapped_key = base64.b64decode(unwrapped_key)
+
+    # Construct Deidentify Config
+    transformation = {
+        "primitive_transformation": {
+            "crypto_replace_ffx_fpe_config": {
+                "crypto_key": {
+                    "unwrapped": {"key": unwrapped_key}
+                },
+                "common_alphabet": alphabet,
+                "surrogate_info_type": {"name": surrogate_type},
+            }
+        }
+    }
+
+    reidentify_config = {
+        "info_type_transformations": {
+            "transformations": [transformation]
+        }
+    }
+
+    inspect_config = {
+        "custom_info_types": [
+            {"info_type": {"name": surrogate_type}, "surrogate_type": {}}
+        ]
+    }
+
+    # Convert string to item
+    item = {"value": input_str}
+
+    # Call the API
+    response = dlp.reidentify_content(
+        parent,
+        inspect_config=inspect_config,
+        reidentify_config=reidentify_config,
+        item=item,
+    )
+
+    # Print results
+    print(response.item.value)
+
+
+# [END dlp_reidentify_free_text_with_fpe_using_surrogate]
+
 
 # [START dlp_deidentify_date_shift]
 def deidentify_with_date_shift(

dlp/deid_test.py

@@ -23,6 +23,7 @@
 HARMFUL_STRING = "My SSN is 372819127"
 HARMLESS_STRING = "My favorite color is blue"
 GCLOUD_PROJECT = os.getenv("GOOGLE_CLOUD_PROJECT")
+UNWRAPPED_KEY = "YWJjZGVmZ2hpamtsbW5vcA=="
 WRAPPED_KEY = (
     "CiQAz0hX4+go8fJwn80Fr8pVImwx+tmZdqU7JL+7TN/S5JxBU9gSSQDhFHpFVy"
     "uzJps0YH9ls480mU+JLG7jI/0lL04i6XJRWqmI6gUSZRUtECYcLH5gXK4SXHlL"
@@ -205,6 +206,24 @@ def test_reidentify_with_fpe(capsys):
     assert "731997681" not in out
 
 
+def test_reidentify_free_text_with_fpe_using_surrogate(capsys):
+    labeled_fpe_string = "My phone number is PHONE_TOKEN(10):9617256398"
+
+    deid.reidentify_free_text_with_fpe_using_surrogate(
+        GCLOUD_PROJECT,
+        labeled_fpe_string,
+        surrogate_type="PHONE_TOKEN",
+        unwrapped_key=UNWRAPPED_KEY,
+        alphabet="NUMERIC",
+    )
+
+    out, _ = capsys.readouterr()
+
+    assert "PHONE_TOKEN" not in out
+    assert "9617256398" not in out
+    assert "My phone number is" in out
+
+
 def test_deidentify_with_replace_infotype(capsys):
     url_to_redact = "https://cloud.google.com"
     deid.deidentify_with_replace_infotype(