[SEC-6587] Databricks CLI Tool Config File inherits default system umask (#522)

Author: shreyas-goenka

databricks_cli/configure/provider.py

@@ -79,9 +79,17 @@ def _set_option(raw_config, profile, option, value):
 
 def _overwrite_config(raw_config):
     config_path = _get_path()
+    # Create config file with owner only rw permissions
+    if not os.path.exists(config_path):
+        file_descriptor = os.open(config_path, os.O_CREAT | os.O_RDWR, 0o600)
+        os.close(file_descriptor)
+
+    # Change file permissions to owner only rw if that's not the case
+    if not os.stat(config_path).st_mode == 0o100600:
+        os.chmod(config_path, 0o600)
+
     with open(config_path, 'w') as cfg:
         raw_config.write(cfg)
-    os.chmod(config_path, 0o600)
 
 
 def update_and_persist_config(profile, databricks_config):

tests/configure/test_provider.py

@@ -23,16 +23,19 @@
 
 import os
 
+from configparser import ConfigParser
 from mock import patch
 import pytest
 
 from databricks_cli.configure.provider import DatabricksConfig, DEFAULT_SECTION, \
     update_and_persist_config, get_config_for_profile, get_config, \
     set_config_provider, ProfileConfigProvider, _get_path, DatabricksConfigProvider,\
-    SparkTaskContextConfigProvider
+    SparkTaskContextConfigProvider, _overwrite_config
 from databricks_cli.utils import InvalidConfigurationError
 
 
+
+
 TEST_HOST = 'https://test.cloud.databricks.com'
 TEST_USER = '[email protected]'
 TEST_PASSWORD = 'banana' # NOQA
@@ -246,3 +249,25 @@ def test_mlflow_config_constructor():
     assert conf.password == TEST_PASSWORD
     assert conf.token == TEST_TOKEN
     assert conf.insecure is False
+
+def test_overwrite_config_creates_file_with_correct_permission():
+    config_path = _get_path()
+
+    assert not os.path.exists(config_path)
+    _overwrite_config(ConfigParser())
+    assert os.path.exists(config_path)
+
+    # assert mode 600 ie owner only can read write
+    assert os.stat(config_path).st_mode == 0o100600
+
+
+def test_overwrite_config_overwrites_permissions_to_600():
+    config_path = _get_path()
+    file_descriptor = os.open(config_path, os.O_CREAT | os.O_RDWR)
+    os.close(file_descriptor)
+
+    assert not os.stat(config_path).st_mode == 0o100600
+
+    _overwrite_config(ConfigParser())
+
+    assert os.stat(config_path).st_mode == 0o100600