From 0174af9f91c06126629a0d4a6f86b9d687cd0a9b Mon Sep 17 00:00:00 2001
From: Georgi <georgi.ajaeiya@datadoghq.com>
Date: Fri, 7 Mar 2025 08:45:02 +0100
Subject: [PATCH] [AWSX][Logs Forwarder] Enable Email and IP redaction based on
 user config

---
 aws/logs_monitoring/logs/datadog_scrubber.py |  9 ++++--
 aws/logs_monitoring/settings.py              |  9 ++++--
 aws/logs_monitoring/tests/test_logs.py       | 31 ++++++++++++--------
 3 files changed, 32 insertions(+), 17 deletions(-)

diff --git a/aws/logs_monitoring/logs/datadog_scrubber.py b/aws/logs_monitoring/logs/datadog_scrubber.py
index 622b21c37..f4f756e83 100644
--- a/aws/logs_monitoring/logs/datadog_scrubber.py
+++ b/aws/logs_monitoring/logs/datadog_scrubber.py
@@ -16,13 +16,17 @@ def __init__(self, configs):
             if config.name in os.environ:
                 rules.append(
                     ScrubbingRule(
-                        compileRegex(config.name, config.pattern), config.placeholder
+                        compileRegex(config.name, config.pattern),
+                        config.placeholder,
+                        config.enabled,
                     )
                 )
         self._rules = rules
 
     def scrub(self, payload):
         for rule in self._rules:
+            if rule.enabled is False:
+                continue
             try:
                 payload = rule.regex.sub(rule.placeholder, payload)
             except Exception:
@@ -31,6 +35,7 @@ def scrub(self, payload):
 
 
 class ScrubbingRule(object):
-    def __init__(self, regex, placeholder):
+    def __init__(self, regex, placeholder, enabled):
         self.regex = regex
         self.placeholder = placeholder
+        self.enabled = enabled
diff --git a/aws/logs_monitoring/settings.py b/aws/logs_monitoring/settings.py
index dd5ca2855..d0fa4aed4 100644
--- a/aws/logs_monitoring/settings.py
+++ b/aws/logs_monitoring/settings.py
@@ -156,22 +156,27 @@ def get_env_var(envvar, default, boolean=False):
 
 
 class ScrubbingRuleConfig(object):
-    def __init__(self, name, pattern, placeholder):
+    def __init__(self, name, pattern, placeholder, enabled=True):
         self.name = name
         self.pattern = pattern
         self.placeholder = placeholder
+        self.enabled = enabled
 
 
 # Scrubbing sensitive data
 # Option to redact all pattern that looks like an ip address / email address / custom pattern
 SCRUBBING_RULE_CONFIGS = [
     ScrubbingRuleConfig(
-        "REDACT_IP", r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", "xxx.xxx.xxx.xxx"
+        "REDACT_IP",
+        r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}",
+        "xxx.xxx.xxx.xxx",
+        get_env_var("REDACT_IP", "false", boolean=True),
     ),
     ScrubbingRuleConfig(
         "REDACT_EMAIL",
         r"[a-zA-Z0-9_.+-]+@[a-zA-Z0-9-]+\.[a-zA-Z0-9-.]+",
         "xxxxx@xxxxx.com",
+        get_env_var("REDACT_EMAIL", "false", boolean=True),
     ),
     ScrubbingRuleConfig(
         "DD_SCRUBBING_RULE",
diff --git a/aws/logs_monitoring/tests/test_logs.py b/aws/logs_monitoring/tests/test_logs.py
index b373ee863..d245b4af7 100644
--- a/aws/logs_monitoring/tests/test_logs.py
+++ b/aws/logs_monitoring/tests/test_logs.py
@@ -1,16 +1,20 @@
 import unittest
 import os
+import sys
+from importlib import reload
+import unittest.mock
 
 from logs.datadog_scrubber import DatadogScrubber
 from logs.datadog_batcher import DatadogBatcher
 from logs.helpers import filter_logs
-from settings import ScrubbingRuleConfig, SCRUBBING_RULE_CONFIGS, get_env_var
 
 
 class TestScrubLogs(unittest.TestCase):
+    @unittest.mock.patch.dict(os.environ, {"REDACT_IP": "true", "REDACT_EMAIL": "true"})
     def test_scrubbing_rule_config(self):
-        os.environ["REDACT_IP"] = ""
-        os.environ["REDACT_EMAIL"] = ""
+        reload(sys.modules["settings"])
+        from settings import SCRUBBING_RULE_CONFIGS
+
         scrubber = DatadogScrubber(SCRUBBING_RULE_CONFIGS)
         payload = scrubber.scrub(
             "ip_address is 127.0.0.1, email is abc.edf@example.com"
@@ -21,17 +25,18 @@ def test_scrubbing_rule_config(self):
         os.environ.pop("REDACT_IP", None)
         os.environ.pop("REDACT_EMAIL", None)
 
+    @unittest.mock.patch.dict(
+        os.environ,
+        {
+            "DD_SCRUBBING_RULE": "[^\u0001-\u007f]+",
+            "DD_SCRUBBING_RULE_REPLACEMENT": "xxxxx",
+        },
+    )
     def test_non_ascii(self):
-        os.environ["DD_SCRUBBING_RULE"] = "[^\u0001-\u007f]+"
-        scrubber = DatadogScrubber(
-            [
-                ScrubbingRuleConfig(
-                    "DD_SCRUBBING_RULE",
-                    get_env_var("DD_SCRUBBING_RULE", default=None),
-                    get_env_var("DD_SCRUBBING_RULE_REPLACEMENT", default="xxxxx"),
-                )
-            ]
-        )
+        reload(sys.modules["settings"])
+        from settings import SCRUBBING_RULE_CONFIGS
+
+        scrubber = DatadogScrubber(SCRUBBING_RULE_CONFIGS)
         payload = scrubber.scrub("abcdef日本語efgかきくけこhij")
         self.assertEqual(payload, "abcdefxxxxxefgxxxxxhij")
         os.environ.pop("DD_SCRUBBING_RULE", None)