Handle case-insensitive sensitive headers (#674)

* Handle case-insensitive sensitive headers * Add test cases for hide sensitive headers functionality * Refactor encoded_headers method --------- Co-authored-by: Alexey Nikitin <alexeynikitin@yandex-team.ru>
jazzband · Sep 11, 2023 · 9c2aba2 · 9c2aba2
1 parent 25b3191
commit 9c2aba2
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 8 deletions.
diff --git a/project/tests/test_sensitive_data_in_request.py b/project/tests/test_sensitive_data_in_request.py
@@ -9,6 +9,7 @@
 HTTP_CONTENT_TYPE = 'content-type'
 CLEANSED = RequestModelFactory.CLEANSED_SUBSTITUTE
 DEFAULT_SENSITIVE_KEYS = {'username', 'api', 'token', 'key', 'secret', 'password', 'signature'}
+DEFAULT_HIDE_COOKIES = True
 
 
 class MaskCredentialsInFormsTest(TestCase):
@@ -127,6 +128,10 @@ class TestEncodingForRequests(TestCase):
     Check that the RequestModelFactory masks sensitive data
     """
 
+    def tearDown(self):
+        SilkyConfig().SILKY_SENSITIVE_KEYS = DEFAULT_SENSITIVE_KEYS
+        SilkyConfig().SILKY_HIDE_COOKIES = DEFAULT_HIDE_COOKIES
+
     def test_password_in_body(self):
         mock_request = Mock()
         mock_request.headers = {HTTP_CONTENT_TYPE: 'text/plain'}
@@ -200,3 +205,46 @@ def test_authorization_header(self):
 
         self.assertIn('authorization', json_headers)
         self.assertEqual(json_headers['authorization'], RequestModelFactory.CLEANSED_SUBSTITUTE)
+
+    def test_hide_cookies(self):
+        SilkyConfig().SILKY_HIDE_COOKIES = True
+        mock_request = Mock()
+        mock_request.headers = {'Cookie': 'secret'}
+        mock_request.body = ''
+        mock_request.get = mock_request.headers.get
+        factory = RequestModelFactory(mock_request)
+        headers = factory.encoded_headers()
+        json_headers = json.loads(headers)
+
+        self.assertIn('cookie', json_headers)
+        self.assertEqual(json_headers['cookie'], RequestModelFactory.CLEANSED_SUBSTITUTE)
+
+    def test_no_hide_cookies(self):
+        SilkyConfig().SILKY_HIDE_COOKIES = False
+        mock_request = Mock()
+        mock_request.headers = {'Cookie': 'Cookies!!!'}
+        mock_request.body = ''
+        mock_request.get = mock_request.headers.get
+        factory = RequestModelFactory(mock_request)
+        headers = factory.encoded_headers()
+        json_headers = json.loads(headers)
+
+        self.assertIn('cookie', json_headers)
+        self.assertEqual(json_headers['cookie'], 'Cookies!!!')
+
+    def test_hide_sensitive_headers(self):
+        SilkyConfig().SILKY_SENSITIVE_KEYS = ["foo", "bar"]
+        mock_request = Mock()
+        mock_request.headers = {'FOO': 'secret', 'BAR': 'secret', 'BAZ': 'not-secret'}
+        mock_request.body = ''
+        mock_request.get = mock_request.headers.get
+        factory = RequestModelFactory(mock_request)
+        headers = factory.encoded_headers()
+        json_headers = json.loads(headers)
+
+        self.assertIn('foo', json_headers)
+        self.assertIn('bar', json_headers)
+        self.assertIn('baz', json_headers)
+        self.assertEqual(json_headers['foo'], RequestModelFactory.CLEANSED_SUBSTITUTE)
+        self.assertEqual(json_headers['bar'], RequestModelFactory.CLEANSED_SUBSTITUTE)
+        self.assertEqual(json_headers['baz'], 'not-secret')
diff --git a/silk/model_factory.py b/silk/model_factory.py
@@ -66,19 +66,17 @@ def encoded_headers(self):
         """
         From Django docs (https://docs.djangoproject.com/en/2.0/ref/request-response/#httprequest-objects):
         """
-        headers = {}
-        sensitive_headers = {'authorization'}
+        sensitive_headers = set(map(str.lower, SilkyConfig().SILKY_SENSITIVE_KEYS))
+        sensitive_headers.add('authorization')
+        if SilkyConfig().SILKY_HIDE_COOKIES:
+            sensitive_headers.add('cookie')
 
+        headers = {}
         for k, v in self.request.headers.items():
+            k = k.lower()
             if k in sensitive_headers:
                 v = RequestModelFactory.CLEANSED_SUBSTITUTE
-
             headers[k] = v
-        if SilkyConfig().SILKY_HIDE_COOKIES:
-            try:
-                del headers['COOKIE']
-            except KeyError:
-                pass
 
         return json.dumps(headers, cls=DefaultEncoder, ensure_ascii=SilkyConfig().SILKY_JSON_ENSURE_ASCII)