Refactor with methods on AgentCheck

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
 
+    @property
+    def _sanitizer(self):
+        # type: () -> SecretsSanitizer
+        if not hasattr(self, '_secrets_sanitizer'):
+            # Configure lazily so that checks that don't use sanitization aren't affected.
+            self._secrets_sanitizer = SecretsSanitizer()
+            self.log.logger.addFilter(SanitizationFilter('secrets', sanitize=self.sanitize))
+
+        return self._secrets_sanitizer
+
+    def register_secret(self, secret):
+        # type: (str) -> None
+        """
+        Register a secret to be scrubbed by `.sanitize()`.
+        """
+        self._sanitizer.register(secret)
+
+    def sanitize(self, text):
+        # type: (str) -> str
+        """
+        Scrub any registered secrets in `text`.
+        """
+        return self._sanitizer.sanitize(text)
+
     def get_instance_proxy(self, instance, uri, proxies=None):
         # type: (InstanceType, str, ProxySettings) -> ProxySettings
         # TODO: Remove with Agent 5
@@ -563,8 +589,6 @@ def service_check(self, name, status, tags=None, hostname=None, message=None, ra
         else:
             message = to_native_string(message)
 
-        message = self.log.redact_registered_secrets(message)
-
         aggregator.submit_service_check(
             self, self.check_id, self._format_namespace(name, raw), status, tags, hostname, message
         )

datadog_checks_base/datadog_checks/base/log.py

@@ -2,7 +2,7 @@
 # All rights reserved
 # Licensed under a 3-clause BSD style license (see LICENSE)
 import logging
-from typing import Any, Set
+from typing import Callable
 
 from six import PY2, iteritems, text_type
 
@@ -24,45 +24,11 @@ def trace(self, msg, *args, **kwargs):
             self._log(TRACE_LEVEL, msg, args, **kwargs)
 
 
-class RedactFilter(logging.Filter):
-    def __init__(self, *args, **kwargs):
-        # type: (*Any, **Any) -> None
-        super(RedactFilter, self).__init__(*args, **kwargs)
-        self.patterns = set()  # type: Set[str]
-
-    def filter(self, record):
-        # type: (logging.LogRecord) -> bool
-        record.msg = self.redact(to_native_string(record.msg))
-
-        if isinstance(record.args, dict):
-            record.args = {key: self.redact(value) for key, value in iteritems(record.args)}
-        else:
-            record.args = tuple(self.redact(arg) for arg in record.args)
-
-        return True
-
-    def redact(self, message):
-        # type: (str) -> str
-        for pattern in self.patterns:
-            message = message.replace(pattern, '********')
-        return message
-
-
 class CheckLoggingAdapter(logging.LoggerAdapter):
     def __init__(self, logger, check):
         super(CheckLoggingAdapter, self).__init__(logger, {})
         self.check = check
         self.check_id = self.check.check_id
-        self._redact_filter = RedactFilter('redact-secrets')
-        self.logger.addFilter(self._redact_filter)
-
-    def register_secret_for_redaction(self, secret):
-        # type: (str) -> None
-        self._redact_filter.patterns.add(secret)
-
-    def redact_registered_secrets(self, message):
-        # type: (str) -> str
-        return self._redact_filter.redact(message)
 
     def process(self, msg, kwargs):
         # Cache for performance
@@ -101,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
@@ -14,6 +15,7 @@
 from datadog_checks.base import __version__ as base_package_version
 from datadog_checks.base import to_native_string
 from datadog_checks.base.checks.base import datadog_agent
+from datadog_checks.base.utils.secrets import SecretsSanitizer
 
 
 def test_instance():
@@ -65,27 +67,53 @@ def test_log_critical_error():
         check.log.critical('test')
 
 
-class TestLogSecretRedaction:
-    def test_default(self, caplog):
-        # By default, non-registered secrets aren't redacted.
-        check = AgentCheck()
+class TestSecretsSanitization:
+    def test_sanitizer(self):
+        # type: () -> None
         secret = 'p@$$w0rd'
-        check.log.error('hello, %s', secret)
-        assert secret in caplog.text
+        sanitizer = SecretsSanitizer()
+        assert sanitizer.sanitize(secret) == secret
+        sanitizer.register(secret)
+        sanitized = sanitizer.sanitize(secret)
+        assert all(letter == '*' for letter in sanitized)
 
-    def test_redact_logs(self, caplog):
+    def test_default(self):
+        # type: () -> None
+        secret = 'p@$$w0rd'
         check = AgentCheck()
+        assert check.sanitize(secret) == secret
+
+    def test_sanitize_text(self):
+        # type: () -> None
         secret = 'p@$$w0rd'
-        check.log.register_secret_for_redaction(secret)
+        check = AgentCheck()
+        check.register_secret(secret)
+
+        sanitized = check.sanitize(secret)
+        assert sanitized != ''
+        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_redact_service_check_messages(self, aggregator, caplog):
-        check = AgentCheck()
+    def test_sanitize_service_checks(self, aggregator, caplog):
+        # type: (Any, Any) -> None
+        # NOTE: no special logic being tested here, but this showcases a possible use case.
         secret = 'p@$$w0rd'
-        check.log.register_secret_for_redaction(secret)
-        check.service_check('test.can_check', status=AgentCheck.CRITICAL, message=secret)
-        aggregator.assert_service_check('test.can_check', status=AgentCheck.CRITICAL, message='********')
+        check = AgentCheck()
+        check.register_secret(secret)
+
+        sanitized = check.sanitize(secret)
+        assert secret not in sanitized
+        check.service_check('test.can_check', status=AgentCheck.CRITICAL, message=sanitized)
+
+        aggregator.assert_service_check('test.can_check', status=AgentCheck.CRITICAL, message=sanitized)
 
 
 def test_warning_ok():