Add secrets sanitization helpers

datadog_checks_base/datadog_checks/base/checks/base.py

@@ -33,6 +33,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 +338,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.setup_sanitization(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 +588,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 +800,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,6 +2,7 @@
 # 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
 
@@ -29,6 +30,12 @@ def __init__(self, logger, check):
         self.check = check
         self.check_id = self.check.check_id
 
+    def setup_sanitization(self, sanitize):
+        # type: (Callable[[str], str]) -> None
+        for handler in self.logger.handlers:
+            if isinstance(handler, AgentLogHandler):
+                handler.setFormatter(SanitizationFormatter(handler.formatter, sanitize=sanitize))
+
     def process(self, msg, kwargs):
         # Cache for performance
         if not self.check_id:
@@ -49,21 +56,49 @@ def warn(self, msg, *args, **kwargs):
             self.log(logging.WARNING, msg, *args, **kwargs)
 
 
+class CheckLogFormatter(logging.Formatter):
+    def format(self, record):
+        # type: (logging.LogRecord) -> str
+        message = to_native_string(super(CheckLogFormatter, self).format(record))
+        return "{} | ({}:{}) | {}".format(
+            # Default to `-` for non-check logs
+            getattr(record, '_check_id', '-'),
+            getattr(record, '_filename', record.filename),
+            getattr(record, '_lineno', record.lineno),
+            message,
+        )
+
+
 class AgentLogHandler(logging.Handler):
     """
     This handler forwards every log to the Go backend allowing python checks to
     log message within the main agent logging system.
     """
 
+    def __init__(self):
+        # type: () -> None
+        super(AgentLogHandler, self).__init__()
+        self.formatter = CheckLogFormatter()  # type: logging.Formatter
+
     def emit(self, record):
-        msg = "{} | ({}:{}) | {}".format(
-            # Default to `-` for non-check logs
-            getattr(record, '_check_id', '-'),
-            getattr(record, '_filename', record.filename),
-            getattr(record, '_lineno', record.lineno),
-            to_native_string(self.format(record)),
-        )
-        datadog_agent.log(msg, record.levelno)
+        # type: (logging.LogRecord) -> None
+        message = self.format(record)
+        datadog_agent.log(message, record.levelno)
+
+
+class SanitizationFormatter(logging.Formatter):
+    """
+    A formatter-like object that sanitizes log messages to hide sensitive data.
+    """
+
+    def __init__(self, parent, sanitize):
+        # type: (logging.Formatter, Callable[[str], str]) -> None
+        self.parent = parent
+        self.sanitize = sanitize
+
+    def format(self, record):
+        # type: (logging.LogRecord) -> str
+        return self.sanitize(self.parent.format(record))
 
 
 LOG_LEVEL_MAP = {
@@ -98,6 +133,7 @@ def _get_py_loglevel(lvl):
 
 
 def init_logging():
+    # type: () -> None
     """
     Initialize logging (set up forwarding to Go backend and sane defaults)
     """

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