Merge pull request #474 from rollbar/fixed/sdk-426-scrubbing-for-namedtuples

Fixed SDK-426 scrubbing for namedtuples

Author: danielmorell

rollbar/__init__.py

@@ -333,6 +333,7 @@ def _get_fastapi_request():
 # Set in init()
 _transforms = []
 _serialize_transform = None
+_scrub_redact_transform = None
 
 _initialized = False
 
@@ -362,7 +363,7 @@ def init(access_token, environment='production', scrub_fields=None, url_fields=N
                  'staging', 'yourname'
     **kw: provided keyword arguments will override keys in SETTINGS.
     """
-    global SETTINGS, agent_log, _initialized, _transforms, _serialize_transform, _threads
+    global SETTINGS, agent_log, _initialized, _transforms, _serialize_transform, _scrub_redact_transform, _threads
 
     if scrub_fields is not None:
        SETTINGS['scrub_fields'] = list(scrub_fields)
@@ -405,6 +406,8 @@ def init(access_token, environment='production', scrub_fields=None, url_fields=N
     _serialize_transform = SerializableTransform(safe_repr=SETTINGS['locals']['safe_repr'],
                                                  safelist_types=SETTINGS['locals']['safelisted_types'])
 
+    _scrub_redact_transform = ScrubRedactTransform(suffixes=[(field,) for field in SETTINGS['scrub_fields']], redact_char='*')
+
     # A list of key prefixes to apply our shortener transform to. The request
     # being included in the body key is old behavior and is being retained for
     # backwards compatibility.
@@ -429,13 +432,10 @@ def init(access_token, environment='production', scrub_fields=None, url_fields=N
                                    **SETTINGS['locals']['sizes'])
     _transforms = [
         shortener,  # priority: 10
-        ScrubRedactTransform(),  # priority: 20
+        _scrub_redact_transform,  # priority: 20
         _serialize_transform,  # priority: 30
-        ScrubTransform(suffixes=[(field,) for field in SETTINGS['scrub_fields']], redact_char='*'),  # priority: 40
-        ScrubUrlTransform(
-            suffixes=[(field,) for field in SETTINGS['url_fields']],
-            params_to_scrub=SETTINGS['scrub_fields'],
-        )  # priority: 50
+        ScrubUrlTransform(suffixes=[(field,) for field in SETTINGS['url_fields']],
+                          params_to_scrub=SETTINGS['scrub_fields'])  # priority: 50
     ]
 
     # Add custom transforms
@@ -1079,7 +1079,7 @@ def _add_locals_data(trace_data, exc_info):
 def _serialize_frame_data(data):
     return transforms.transform(
         data,
-        [ScrubRedactTransform(), _serialize_transform],
+        [_scrub_redact_transform, _serialize_transform],
         batch_transforms=SETTINGS['batch_transforms']
     )
 

rollbar/test/test_rollbar.py

@@ -6,6 +6,7 @@
 import uuid
 
 import sys
+from collections import namedtuple
 
 try:
     from StringIO import StringIO
@@ -1464,6 +1465,29 @@ def _raise():
         undecodable_message = '<Undecodable type:(%s) base64:(%s)>' % ('bytes', base64.b64encode(invalid).decode('ascii'))
         self.assertEqual(undecodable_message, payload['data']['body']['trace']['frames'][-1]['locals']['_invalid'])
 
+    @mock.patch('rollbar.send_payload')
+    def test_scrub_namedtuple(self, send_payload):
+
+        SomeTuple = namedtuple('SomeTuple', ['password', 'some_field'])
+
+        def _raise():
+            Data = SomeTuple(password='clear_text', some_field='some_field')
+
+            password = 'sensitive'
+            raise Exception((Data, password))
+
+        try:
+            _raise()
+        except:
+            rollbar.report_exc_info()
+
+        self.assertEqual(send_payload.called, True)
+
+        payload = send_payload.call_args[0][0]
+
+        self.assertRegex(payload['data']['body']['trace']['frames'][-1]['locals']['password'], r'\*+')
+        self.assertRegex(payload['data']['body']['trace']['frames'][-1]['locals']['Data'], 'password=\'\*+\'')
+
     @mock.patch('rollbar.send_payload')
     def test_scrub_nans(self, send_payload):
         def _raise():

rollbar/test/test_scrub_redact_transform.py

@@ -58,12 +58,19 @@ def test_scrub_set(self):
         expected = set([SCRUBBED, NOT_REDACT_REF])
         self._assertScrubbed(obj, expected)
 
-    def scrub_tuple(self):
+    def test_scrub_tuple(self):
         obj = (REDACT_REF, REDACT_REF, REDACT_REF)
         expected = (SCRUBBED, SCRUBBED, SCRUBBED)
         self._assertScrubbed(obj, expected)
 
+    def test_scrub_namedtuple(self):
+        from collections import namedtuple
+        TestNamedTuple = namedtuple('TestNamedTuple', ['scrub_me', 'dont_scrub_me'])
+        obj = TestNamedTuple(scrub_me=REDACT_REF, dont_scrub_me=NOT_REDACT_REF)
+        expected = TestNamedTuple(scrub_me=SCRUBBED, dont_scrub_me=NOT_REDACT_REF)
+        self._assertScrubbed(obj, expected)
+
     def test_scrub_dict(self):
-        obj = {'scrub_me': REDACT_REF}
-        expected = {'scrub_me': SCRUBBED}
+        obj = {'scrub_me': REDACT_REF, 'dont_scrub_me': NOT_REDACT_REF}
+        expected = {'scrub_me': SCRUBBED, 'dont_scrub_me': NOT_REDACT_REF}
         self._assertScrubbed(obj, expected)

rollbar/test/test_scrub_transform.py

@@ -61,6 +61,20 @@ def test_no_suffixes(self):
         expected = dict(obj)
         self._assertScrubbed([], obj, expected)
 
+    def test_scrub_named_tuple(self):
+        from collections import namedtuple
+        Obj = namedtuple('obj', ['hello', 'password'])
+        obj = Obj('world', 'cleartext')
+        expected = Obj('world', '*********')
+        self._assertScrubbed([['password']], obj, expected)
+
+    def test_scrub_named_tuple_in_list(self):
+        from collections import namedtuple
+        Obj = namedtuple('obj', ['hello', 'password'])
+        obj = [Obj('world', 'cleartext'), 'another element']
+        expected = [Obj('world', '*********'), 'another element']
+        self._assertScrubbed([['password']], obj, expected)
+
     def test_scrub_simple_dict(self):
         obj = {'hello': 'world', 'password': 'cleartext'}
         expected = copy.deepcopy(obj)

rollbar/test/test_transform.py

@@ -26,7 +26,6 @@ def test_default_transforms(self):
             ShortenerTransform,
             ScrubRedactTransform,
             SerializableTransform,
-            ScrubTransform,
             ScrubUrlTransform,
         }, transforms)
 
@@ -49,7 +48,6 @@ class CustomTransform(Transform):
             ShortenerTransform,
             ScrubRedactTransform,
             SerializableTransform,
-            ScrubTransform,
             ScrubUrlTransform,
             CustomTransform,
         }, transforms)
@@ -59,7 +57,6 @@ class CustomTransform(Transform):
             ShortenerTransform,
             ScrubRedactTransform,
             SerializableTransform,
-            ScrubTransform,
             ScrubUrlTransform,
             CustomTransform,
         ], transforms_ordered)
@@ -88,6 +85,5 @@ class CustomTransformTwo(Transform):
             ScrubRedactTransform,  # priority 20
             SerializableTransform,  # priority 30
             CustomTransformTwo,  # priority 35
-            ScrubTransform,  # priority 40
             ScrubUrlTransform,  # priority 50
         ], transforms)