[xy] Mask env var values with stars in terminal output (mage-ai#2201)

* [xy] Hide environment variables.

* [xy] Add helper file to filter out env var values.

* [xy] Add HIDE_ENV_VAR_VALUES variable.

* [xy] Add unit test.

mage_ai/server/websocket_server.py

@@ -39,9 +39,11 @@
 )
 from mage_ai.settings import (
     DISABLE_NOTEBOOK_EDIT_ACCESS,
+    HIDE_ENV_VAR_VALUES,
     REQUIRE_USER_AUTHENTICATION,
 )
 from mage_ai.shared.hash import merge_dict
+from mage_ai.shared.security import filter_out_env_var_values
 from mage_ai.utils.code import reload_all_repo_modules
 from jupyter_client import KernelClient
 from typing import Dict
@@ -259,6 +261,16 @@ def should_filter_message(message):
 
             return False
 
+        def filter_out_sensitive_data(message):
+            if not message.get('data') or not HIDE_ENV_VAR_VALUES:
+                return message
+            data = message['data']
+            if type(data) is str:
+                data = [data]
+            data = [filter_out_env_var_values(data_value) for data_value in data]
+            message['data'] = data
+            return message
+
         msg_id = message.get('msg_id')
         if msg_id is None:
             return
@@ -269,6 +281,8 @@ def should_filter_message(message):
         if should_filter_message(message):
             return
 
+        message = filter_out_sensitive_data(message)
+
         execution_metadata = message.get('execution_metadata')
         msg_id_value = execution_metadata if execution_metadata is not None \
             else WebSocketServer.running_executions_mapping.get(msg_id, dict())

mage_ai/settings/__init__.py

@@ -2,6 +2,7 @@
 
 
 DEBUG = os.getenv('DEBUG', False)
+HIDE_ENV_VAR_VALUES = int(os.getenv('HIDE_ENV_VAR_VALUES', 1) or 1) == 1
 QUERY_API_KEY = 'api_key'
 
 """

mage_ai/shared/security.py

@@ -0,0 +1,14 @@
+import os
+
+MIN_SECRET_ENV_VAR_LENGTH = 5
+
+
+def filter_out_env_var_values(value: str):
+    env_var_values = dict(os.environ).values()
+    env_var_values = [v for v in env_var_values if v and len(v) >= MIN_SECRET_ENV_VAR_LENGTH]
+    env_var_values.sort(key=len, reverse=True)
+    value_clean = value
+    for env_var_value in env_var_values:
+        replace_value = '*' * len(env_var_value)
+        value_clean = value_clean.replace(env_var_value, replace_value)
+    return value_clean

mage_ai/tests/shared/test_security.py

@@ -0,0 +1,23 @@
+from mage_ai.shared.security import filter_out_env_var_values
+from mage_ai.tests.base_test import TestCase
+from unittest.mock import patch
+import os
+
+
+MOCK_ENV_VARS = {
+    'VAR1': '123',
+    'VAR2': '45678',
+    'VAR3': 'abcdefg',
+    'VAR4': '',
+}
+
+
+class SecurityTests(TestCase):
+    @patch.dict(os.environ, MOCK_ENV_VARS)
+    def test_filter_out_env_var_values(self):
+        value1 = filter_out_env_var_values('12345678abcdefghij')
+        value2 = filter_out_env_var_values('testdata')
+        value3 = filter_out_env_var_values('test45645678')
+        self.assertEqual(value1, '123************hij')
+        self.assertEqual(value2, 'testdata')
+        self.assertEqual(value3, 'test456*****')