Add secrets sanitization helpers

datadog_checks_base/datadog_checks/base/checks/base.py

@@ -18,6 +18,7 @@
 
 from ..config import is_affirmative
 from ..constants import ServiceCheck
+from ..log import SanitizationFilter
 from ..types import (
     AgentConfigType,
     Event,
@@ -33,6 +34,7 @@
 from ..utils.limiter import Limiter
 from ..utils.metadata import MetadataManager
 from ..utils.proxy import config_proxy_skip
+from ..utils.secrets import SecretsSanitizer
 
 try:
     import datadog_agent
@@ -337,6 +339,30 @@ def in_developer_mode(self):
         self._log_deprecation('in_developer_mode')
         return False
 
+    def register_secret(self, secret):
+        # type: (str) -> None
+        """
+        Register a secret to be scrubbed by `.sanitize()`.
+        """
+        if not hasattr(self, '_sanitizer'):
+            # Configure lazily so that checks that don't use sanitization aren't affected.
+            self._sanitizer = SecretsSanitizer()
+            self.log.logger.addFilter(SanitizationFilter('secrets', sanitize=self.sanitize))
+
+        self._sanitizer.register(secret)
+
+    def sanitize(self, text):
+        # type: (str) -> str
+        """
+        Scrub any registered secrets in `text`.
+        """
+        try:
+            sanitizer = self._sanitizer
+        except AttributeError:
+            return text
+        else:
+            return sanitizer.sanitize(text)
+
     def get_instance_proxy(self, instance, uri, proxies=None):
         # type: (InstanceType, str, ProxySettings) -> ProxySettings
         # TODO: Remove with Agent 5
@@ -563,6 +589,8 @@ def service_check(self, name, status, tags=None, hostname=None, message=None, ra
         else:
             message = to_native_string(message)
 
+        message = self.sanitize(message)
+
         aggregator.submit_service_check(
             self, self.check_id, self._format_namespace(name, raw), status, tags, hostname, message
         )
@@ -773,7 +801,9 @@ def run(self):
 
             result = ''
         except Exception as e:
-            result = json.dumps([{'message': str(e), 'traceback': traceback.format_exc()}])
+            message = self.sanitize(str(e))
+            tb = self.sanitize(traceback.format_exc())
+            result = json.dumps([{'message': message, 'traceback': tb}])
         finally:
             if self.metric_limiter:
                 self.metric_limiter.reset()

datadog_checks_base/datadog_checks/base/log.py

@@ -2,8 +2,9 @@
 # All rights reserved
 # Licensed under a 3-clause BSD style license (see LICENSE)
 import logging
+from typing import Callable
 
-from six import PY2, text_type
+from six import PY2, iteritems, text_type
 
 from .utils.common import to_native_string
 
@@ -66,6 +67,28 @@ def emit(self, record):
         datadog_agent.log(msg, record.levelno)
 
 
+class SanitizationFilter(logging.Filter):
+    """
+    A filter for sanitizing log records messages.
+    """
+
+    def __init__(self, name, sanitize):
+        # type: (str, Callable[[str], str]) -> None
+        super(SanitizationFilter, self).__init__(name)
+        self.sanitize = sanitize
+
+    def filter(self, record):
+        # type: (logging.LogRecord) -> bool
+        record.msg = self.sanitize(to_native_string(record.msg))
+
+        if isinstance(record.args, dict):
+            record.args = {key: self.sanitize(value) for key, value in iteritems(record.args)}
+        else:
+            record.args = tuple(self.sanitize(arg) for arg in record.args)
+
+        return True
+
+
 LOG_LEVEL_MAP = {
     'CRIT': logging.CRITICAL,
     'CRITICAL': logging.CRITICAL,

datadog_checks_base/datadog_checks/base/utils/secrets.py

@@ -0,0 +1,23 @@
+from typing import Set
+
+
+class SecretsSanitizer:
+    """
+    A helper for sanitizing secrets (password, keys, ...) in text output.
+    """
+
+    REDACTED = '********'
+
+    def __init__(self):
+        # type: () -> None
+        self.patterns = set()  # type: Set[str]
+
+    def register(self, secret):
+        # type: (str) -> None
+        self.patterns.add(secret)
+
+    def sanitize(self, text):
+        # type: (str) -> str
+        for pattern in self.patterns:
+            text = text.replace(pattern, self.REDACTED)
+        return text

datadog_checks_base/tests/test_agent_check.py

@@ -5,6 +5,7 @@
 # Licensed under a 3-clause BSD style license (see LICENSE)
 import json
 from collections import OrderedDict
+from typing import Any
 
 import mock
 import pytest
@@ -65,6 +66,71 @@ def test_log_critical_error():
         check.log.critical('test')
 
 
+class TestSecretsSanitization:
+    def test_default(self, caplog):
+        # type: (Any) -> None
+        secret = 's3kr3t'
+        check = AgentCheck()
+
+        message = 'hello, {}'.format(secret)
+        assert check.sanitize(message) == message
+
+        check.log.error(message)
+        assert secret in caplog.text
+
+    def test_sanitize_text(self):
+        # type: () -> None
+        secret = 'p@$$w0rd'
+        check = AgentCheck()
+        check.register_secret(secret)
+
+        sanitized = check.sanitize('hello, {}'.format(secret))
+        assert secret not in sanitized
+
+    def text_sanitize_logs(self, caplog):
+        # type: (Any) -> None
+        secret = 'p@$$w0rd'
+        check = AgentCheck()
+        check.register_secret(secret)
+
+        check.log.error('hello, %s', secret)
+        assert secret not in caplog.text
+
+    def test_sanitize_service_check_message(self, aggregator, caplog):
+        # type: (Any, Any) -> None
+        secret = 'p@$$w0rd'
+        check = AgentCheck()
+        check.register_secret(secret)
+        sanitized = check.sanitize(secret)
+
+        check.service_check('test.can_check', status=AgentCheck.CRITICAL, message=secret)
+
+        aggregator.assert_service_check('test.can_check', status=AgentCheck.CRITICAL, message=sanitized)
+
+    def test_sanitize_exception_tracebacks(self):
+        # type: () -> None
+        class MyCheck(AgentCheck):
+            def __init__(self, *args, **kwargs):
+                # type: (*Any, **Any) -> None
+                super(MyCheck, self).__init__(*args, **kwargs)
+                self.password = 'p@$$w0rd'
+                self.register_secret(self.password)
+
+            def check(self, instance):
+                # type: (Any) -> None
+                try:
+                    # Simulate a failing call in a dependency.
+                    raise Exception('Could not establish connection with Password={}'.format(self.password))
+                except Exception as exc:
+                    raise RuntimeError('Unexpected error while executing check: {}'.format(exc))
+
+        check = MyCheck('my_check', {}, [{}])
+        result = json.loads(check.run())[0]
+
+        assert check.password not in result['message']
+        assert check.password not in result['traceback']
+
+
 def test_warning_ok():
     check = AgentCheck()
 

datadog_checks_base/tests/test_utils.py

@@ -20,6 +20,7 @@
 )
 from datadog_checks.base.utils.containers import iter_unique
 from datadog_checks.base.utils.limiter import Limiter
+from datadog_checks.base.utils.secrets import SecretsSanitizer
 
 
 class Item:
@@ -183,3 +184,30 @@ def test_to_string_deprecated(self):
         # type: () -> None
         with pytest.deprecated_call():
             to_string(b'example')
+
+
+class TestSecretsSanitizer:
+    def test_default(self):
+        # type: () -> None
+        secret = 's3kr3t'
+        sanitizer = SecretsSanitizer()
+        assert sanitizer.sanitize(secret) == secret
+
+    def test_sanitize(self):
+        # type: () -> None
+        secret = 's3kr3t'
+        sanitizer = SecretsSanitizer()
+        sanitizer.register(secret)
+        assert all(letter == '*' for letter in sanitizer.sanitize(secret))
+
+    def test_sanitize_multiple(self):
+        # type: () -> None
+        pwd1 = 's3kr3t'
+        pwd2 = 'admin123'
+        sanitizer = SecretsSanitizer()
+        sanitizer.register(pwd1)
+        sanitizer.register(pwd2)
+        message = 'Could not authenticate with password {}, did you try {}?'.format(pwd1, pwd2)
+        sanitized = sanitizer.sanitize(message)
+        assert pwd1 not in sanitized
+        assert pwd2 not in sanitized