From e8aacf959b168124d51937aa8768365557620f86 Mon Sep 17 00:00:00 2001
From: Frode Nordahl <frode.nordahl@canonical.com>
Date: Wed, 9 Jun 2021 10:39:52 +0200
Subject: [PATCH] Write private SSH key with appropriate permissions

Fixes #583
---
 .../test_zaza_utilities_openstack.py          | 21 +++++++++++++++++++
 zaza/openstack/utilities/openstack.py         | 10 +++++++--
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/unit_tests/utilities/test_zaza_utilities_openstack.py b/unit_tests/utilities/test_zaza_utilities_openstack.py
index 3cd33a6..3184166 100644
--- a/unit_tests/utilities/test_zaza_utilities_openstack.py
+++ b/unit_tests/utilities/test_zaza_utilities_openstack.py
@@ -633,15 +633,36 @@ class TestOpenStackUtils(ut_utils.BaseTestCase):
     def test_write_private_key(self):
         self.patch_object(openstack_utils.deployment_env, 'get_tmpdir',
                           return_value='/tmp/zaza-model1')
+        self.patch_object(openstack_utils.os, 'umask',
+                          return_value='fakeumask')
         m = mock.mock_open()
         with mock.patch(
             'zaza.openstack.utilities.openstack.open', m, create=False
         ):
             openstack_utils.write_private_key('mykeys', 'keycontents')
+        self.umask.assert_has_calls([
+            mock.call(0o177),
+            mock.call('fakeumask')
+        ])
         m.assert_called_once_with('/tmp/zaza-model1/id_rsa_mykeys', 'w')
         handle = m()
         handle.write.assert_called_once_with('keycontents')
 
+        # Confirm that umask is reset even if write raises an exception
+        m.reset_mock()
+        self.umask.reset_mock()
+        with mock.patch(
+            'zaza.openstack.utilities.openstack.open', m, create=False
+        ):
+            handle = m()
+            handle.write.side_effect = OSError
+            with self.assertRaises(OSError):
+                openstack_utils.write_private_key('mykeys', 'keycontents')
+        self.umask.assert_has_calls([
+            mock.call(0o177),
+            mock.call('fakeumask')
+        ])
+
     def test_get_private_key(self):
         self.patch_object(openstack_utils.deployment_env, 'get_tmpdir',
                           return_value='/tmp/zaza-model1')
diff --git a/zaza/openstack/utilities/openstack.py b/zaza/openstack/utilities/openstack.py
index 1dfe5a5..9e30359 100644
--- a/zaza/openstack/utilities/openstack.py
+++ b/zaza/openstack/utilities/openstack.py
@@ -2768,8 +2768,14 @@ def write_private_key(keypair_name, key):
     :param key: PEM Encoded Private Key
     :type key: str
     """
-    with open(get_private_key_file(keypair_name), 'w') as key_file:
-        key_file.write(key)
+    # Create the key file with mode 0o600 to allow the developer to pass it to
+    # the `ssh` command without getting a "bad permissions" error.
+    stored_umask = os.umask(0o177)
+    try:
+        with open(get_private_key_file(keypair_name), 'w') as key_file:
+            key_file.write(key)
+    finally:
+        os.umask(stored_umask)
 
 
 def get_private_key(keypair_name):