Addressing -ssh-key-gen bug.

src/azure-cli-core/HISTORY.rst

@@ -5,7 +5,8 @@ Release History
 
 2.0.46
 ++++++
-* Minor fixes.
+* Fixed issue where `az vm create --generate-ssh-keys` overwrites private key
+  file if public key file is missing. (#4725, #6780)
 
 2.0.45
 ++++++

src/azure-cli-core/azure/cli/core/keys.py

@@ -6,6 +6,10 @@
 import os
 import os.path
 
+from knack.util import CLIError
+from knack.log import get_logger
+logger = get_logger(__name__)
+
 
 def is_valid_ssh_rsa_public_key(openssh_pubkey):
     # http://stackoverflow.com/questions/2494450/ssh-rsa-public-key-validation-using-a-regular-expression # pylint: disable=line-too-long
@@ -31,18 +35,44 @@ def is_valid_ssh_rsa_public_key(openssh_pubkey):
 
 def generate_ssh_keys(private_key_filepath, public_key_filepath):
     import paramiko
+    from paramiko.ssh_exception import PasswordRequiredException, SSHException
+
+    if os.path.isfile(public_key_filepath):
+        try:
+            with open(public_key_filepath, 'r') as public_key_file:
+                public_key = public_key_file.read()
+                pub_ssh_dir = os.path.dirname(public_key_filepath)
+                logger.warning("Public SSH key file '%s' found in dir '%s'."
+                               "Use public key file. New RSA key pair will not be generated.",
+                               public_key_filepath, pub_ssh_dir)
 
-    ssh_dir, _ = os.path.split(private_key_filepath)
+                return public_key
+        except IOError as e:
+            raise CLIError(e)
+
+    ssh_dir = os.path.dirname(private_key_filepath)
     if not os.path.exists(ssh_dir):
         os.makedirs(ssh_dir)
         os.chmod(ssh_dir, 0o700)
 
-    key = paramiko.RSAKey.generate(2048)
-    key.write_private_key_file(private_key_filepath)
-    os.chmod(private_key_filepath, 0o600)
+    if os.path.isfile(private_key_filepath):
+        # try to use existing private key if it exists.
+        try:
+            key = paramiko.RSAKey(filename=private_key_filepath)
+            logger.warning("Private SSH key file '%s' found in dir '%s'. "
+                           "Generating new Public key file '%s'",
+                           private_key_filepath, ssh_dir, public_key_filepath)
+        except (PasswordRequiredException, SSHException, IOError) as e:
+            raise CLIError(e)
+
+    else:
+        # otherwise generate new private key.
+        key = paramiko.RSAKey.generate(2048)
+        key.write_private_key_file(private_key_filepath)
+        os.chmod(private_key_filepath, 0o600)
 
     with open(public_key_filepath, 'w') as public_key_file:
-        public_key = '%s %s' % (key.get_name(), key.get_base64())
+        public_key = '{} {}'.format(key.get_name(), key.get_base64())
         public_key_file.write(public_key)
     os.chmod(public_key_filepath, 0o644)
 

src/azure-cli-core/azure/cli/core/tests/test_keys.py

@@ -0,0 +1,133 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+
+import unittest
+import tempfile
+import paramiko
+import io
+import os
+import shutil
+from knack.util import CLIError
+
+from azure.cli.core.keys import generate_ssh_keys
+
+
+class TestGenerateSSHKeys(unittest.TestCase):
+
+    def setUp(self):
+        # set up temporary directory to be used for temp files.
+        self._tempdirName = tempfile.mkdtemp(prefix="key_tmp_")
+
+        self.key = paramiko.RSAKey.generate(2048)
+        keyOutput = io.StringIO()
+        self.key.write_private_key(keyOutput)
+
+        self.private_key = keyOutput.getvalue()
+        self.public_key = '{} {}'.format(self.key.get_name(), self.key.get_base64())
+
+    def tearDown(self):
+        # delete temporary directory to be used for temp files.
+        shutil.rmtree(self._tempdirName)
+
+    def test_when_public_key_file_exists(self):
+        # Create public key file
+        public_key_path = self._create_new_temp_key_file(self.public_key, suffix=".pub")
+
+        # Create private key file
+        private_key_path = self._create_new_temp_key_file(self.private_key)
+
+        # Call generate_ssh_keys and assert that returned public key same as original
+        new_public_key = generate_ssh_keys("", public_key_path)
+        self.assertEqual(self.public_key, new_public_key)
+
+        # Check that private and public key file contents unchanged.
+        with open(public_key_path, 'r') as f:
+            new_public_key = f.read()
+            self.assertEqual(self.public_key, new_public_key)
+
+        with open(private_key_path, 'r') as f:
+            new_private_key = f.read()
+            self.assertEqual(self.private_key, new_private_key)
+
+    def test_error_raised_when_public_key_file_exists_no_permissions(self):
+        # Create public key file with no read or write access
+        public_key_path = self._create_new_temp_key_file(self.public_key)
+        os.chmod(public_key_path, 0o000)
+
+        # Check that CLIError exception is raised when generate_ssh_keys is called.
+        with self.assertRaises(CLIError):
+            generate_ssh_keys("", public_key_path)
+
+    def test_error_raised_when_private_key_file_exists_no_permissions(self):
+        # Create private key file with no read or write access
+        private_key_path = self._create_new_temp_key_file(self.private_key)
+        os.chmod(private_key_path, 0o000)
+
+        # Check that CLIError exception is raised when generate_ssh_keys is called.
+        with self.assertRaises(CLIError):
+            public_key_path = private_key_path + ".pub"
+            generate_ssh_keys(private_key_path, public_key_path)
+
+    def test_error_raised_when_private_key_file_exists_encrypted(self):
+        # Create empty private key file
+        private_key_path = self._create_new_temp_key_file("")
+
+        # Write encrypted key into file
+        self.key.write_private_key_file(private_key_path, "test")
+
+        # Check that CLIError exception is raised when generate_ssh_keys is called.
+        with self.assertRaises(CLIError):
+            public_key_path = private_key_path + ".pub"
+            generate_ssh_keys(private_key_path, public_key_path)
+
+    def test_generate_public_key_file_from_existing_private_key_files(self):
+        # Create private key file
+        private_key_path = self._create_new_temp_key_file(self.private_key)
+
+        # Call generate_ssh_keys and assert that returned public key same as original
+        public_key_path = private_key_path + ".pub"
+        new_public_key = generate_ssh_keys(private_key_path, public_key_path)
+        self.assertEqual(self.public_key, new_public_key)
+
+        # Check that correct public key file has been created
+        with open(public_key_path, 'r') as f:
+            public_key = f.read()
+            self.assertEqual(self.public_key, public_key)
+
+        # Check that private key file contents unchanged
+        with open(private_key_path, 'r') as f:
+            private_key = f.read()
+            self.assertEqual(self.private_key, private_key)
+
+    def test_generate_new_private_public_key_files(self):
+        # create random temp file name
+        f = tempfile.NamedTemporaryFile(mode='w', dir=self._tempdirName)
+        f.close()
+        private_key_path = f.name
+
+        # Call generate_ssh_keys and assert that returned public key same as original
+        public_key_path = private_key_path + ".pub"
+        new_public_key = generate_ssh_keys(private_key_path, public_key_path)
+
+        # Check that public key returned is same as public key in public key path
+        with open(public_key_path, 'r') as f:
+            public_key = f.read()
+            self.assertEqual(public_key, new_public_key)
+
+        # Check that public key corresponds to private key
+        with open(private_key_path, 'r') as f:
+            key = paramiko.RSAKey(filename=private_key_path)
+            public_key = '{} {}'.format(key.get_name(), key.get_base64())
+            self.assertEqual(public_key, new_public_key)
+
+    def _create_new_temp_key_file(self, key_data, suffix=""):
+        with tempfile.NamedTemporaryFile(mode='w', dir=self._tempdirName, delete=False, suffix=suffix) as f:
+            f.write(key_data)
+            return f.name
+
+
+if __name__ == '__main__':
+    unittest.main()

src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/tests/hybrid_2018_03_01/test_vm_actions.py

@@ -4,6 +4,7 @@
 # --------------------------------------------------------------------------------------------
 
 import os
+import shutil
 import tempfile
 import unittest
 import mock
@@ -27,8 +28,16 @@
 
 class TestActions(unittest.TestCase):
     def test_generate_specfied_ssh_key_files(self):
-        _, private_key_file = tempfile.mkstemp()
+        temp_dir_name = tempfile.mkdtemp(prefix="ssh_dir_")
+
+        # cleanup temporary directory and its contents
+        self.addCleanup(shutil.rmtree, path=temp_dir_name)
+
+        # first create file paths for the keys to be generated
+        _, private_key_file = tempfile.mkstemp(dir=temp_dir_name)
         public_key_file = private_key_file + '.pub'
+        os.remove(private_key_file)
+
         args = mock.MagicMock()
         args.ssh_key_value = public_key_file
         args.generate_ssh_keys = True
@@ -51,7 +60,7 @@ def test_generate_specfied_ssh_key_files(self):
         self.assertEqual(generated_public_key_string, args.ssh_key_value)
 
         # 3 verify we do not generate unless told so
-        _, private_key_file2 = tempfile.mkstemp()
+        _, private_key_file2 = tempfile.mkstemp(dir=temp_dir_name)
         public_key_file2 = private_key_file2 + '.pub'
         args3 = mock.MagicMock()
         args3.ssh_key_value = public_key_file2
@@ -60,7 +69,7 @@ def test_generate_specfied_ssh_key_files(self):
             validate_ssh_key(args3)
 
         # 4 verify file naming if the pub file doesn't end with .pub
-        _, public_key_file4 = tempfile.mkstemp()
+        _, public_key_file4 = tempfile.mkstemp(dir=temp_dir_name)
         public_key_file4 += '1'  # make it nonexisting
         args4 = mock.MagicMock()
         args4.ssh_key_value = public_key_file4

src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/tests/latest/test_vm_actions.py

@@ -4,6 +4,7 @@
 # --------------------------------------------------------------------------------------------
 
 import os
+import shutil
 import tempfile
 import unittest
 import mock
@@ -27,8 +28,16 @@
 
 class TestActions(unittest.TestCase):
     def test_generate_specfied_ssh_key_files(self):
-        _, private_key_file = tempfile.mkstemp()
+        temp_dir_name = tempfile.mkdtemp(prefix="ssh_dir_")
+
+        # cleanup temporary directory and its contents
+        self.addCleanup(shutil.rmtree, path=temp_dir_name)
+
+        # first create file paths for the keys to be generated
+        _, private_key_file = tempfile.mkstemp(dir=temp_dir_name)
         public_key_file = private_key_file + '.pub'
+        os.remove(private_key_file)
+
         args = mock.MagicMock()
         args.ssh_key_value = public_key_file
         args.generate_ssh_keys = True
@@ -51,7 +60,7 @@ def test_generate_specfied_ssh_key_files(self):
         self.assertEqual(generated_public_key_string, args.ssh_key_value)
 
         # 3 verify we do not generate unless told so
-        _, private_key_file2 = tempfile.mkstemp()
+        _, private_key_file2 = tempfile.mkstemp(dir=temp_dir_name)
         public_key_file2 = private_key_file2 + '.pub'
         args3 = mock.MagicMock()
         args3.ssh_key_value = public_key_file2
@@ -60,7 +69,7 @@ def test_generate_specfied_ssh_key_files(self):
             validate_ssh_key(args3)
 
         # 4 verify file naming if the pub file doesn't end with .pub
-        _, public_key_file4 = tempfile.mkstemp()
+        _, public_key_file4 = tempfile.mkstemp(dir=temp_dir_name)
         public_key_file4 += '1'  # make it nonexisting
         args4 = mock.MagicMock()
         args4.ssh_key_value = public_key_file4

src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/tests/profile_2017_03_09/test_vm_actions.py

@@ -4,6 +4,7 @@
 # --------------------------------------------------------------------------------------------
 
 import os
+import shutil
 import tempfile
 import unittest
 import mock
@@ -26,8 +27,16 @@
 
 class TestActions(unittest.TestCase):
     def test_generate_specfied_ssh_key_files(self):
-        _, private_key_file = tempfile.mkstemp()
+        temp_dir_name = tempfile.mkdtemp(prefix="ssh_dir_")
+
+        # cleanup temporary directory and its contents
+        self.addCleanup(shutil.rmtree, path=temp_dir_name)
+
+        # first create file paths for the keys to be generated
+        _, private_key_file = tempfile.mkstemp(dir=temp_dir_name)
         public_key_file = private_key_file + '.pub'
+        os.remove(private_key_file)
+
         args = mock.MagicMock()
         args.ssh_key_value = public_key_file
         args.generate_ssh_keys = True
@@ -50,7 +59,7 @@ def test_generate_specfied_ssh_key_files(self):
         self.assertEqual(generated_public_key_string, args.ssh_key_value)
 
         # 3 verify we do not generate unless told so
-        _, private_key_file2 = tempfile.mkstemp()
+        _, private_key_file2 = tempfile.mkstemp(dir=temp_dir_name)
         public_key_file2 = private_key_file2 + '.pub'
         args3 = mock.MagicMock()
         args3.ssh_key_value = public_key_file2
@@ -59,7 +68,7 @@ def test_generate_specfied_ssh_key_files(self):
             validate_ssh_key(args3)
 
         # 4 verify file naming if the pub file doesn't end with .pub
-        _, public_key_file4 = tempfile.mkstemp()
+        _, public_key_file4 = tempfile.mkstemp(dir=temp_dir_name)
         public_key_file4 += '1'  # make it nonexisting
         args4 = mock.MagicMock()
         args4.ssh_key_value = public_key_file4