From 9d43aa3b706057d0fe97fcb41cda353bcee1fb59 Mon Sep 17 00:00:00 2001 From: Oluwatosin Adewale Date: Mon, 20 Aug 2018 16:50:47 -0700 Subject: [PATCH 1/6] Logic to handle situation where "az vm create ... --generate-ssh-keys" is run and id_rsa exists but id_rsa.pub does not exist. --- .../azure/cli/command_modules/vm/_validators.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_validators.py b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_validators.py index e4160d4e34b..c40f64dad61 100644 --- a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_validators.py +++ b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_validators.py @@ -864,15 +864,28 @@ def _validate_admin_password(password, os_type): def validate_ssh_key(namespace): + + # check if id_rsa is present but id_rsa.pub is missing + def _public_key_is_missing(): + private_key_path = os.path.join(os.path.expanduser('~'), '.ssh', 'id_rsa') + public_key_path = os.path.join(os.path.expanduser('~'), '.ssh', 'id_rsa.pub') + return (os.path.exists(private_key_path) and not os.path.exists(public_key_path)) + + + string_or_file = (namespace.ssh_key_value or os.path.join(os.path.expanduser('~'), '.ssh', 'id_rsa.pub')) content = string_or_file if os.path.exists(string_or_file): - logger.info('Use existing SSH public key file: %s', string_or_file) + logger.info('Reusing existing SSH public key from %s', string_or_file) with open(string_or_file, 'r') as f: content = f.read() elif not keys.is_valid_ssh_rsa_public_key(content): if namespace.generate_ssh_keys: + + if _public_key_is_missing(): + raise CLIError('SSH private key id_rsa exists but public key is missing. Please export the public key.') + # figure out appropriate file names: # 'base_name'(with private keys), and 'base_name.pub'(with public keys) public_key_filepath = string_or_file From 7c67fdb0b726d3aae4af5b0a3709ea9a5d4220af Mon Sep 17 00:00:00 2001 From: Oluwatosin Adewale Date: Mon, 20 Aug 2018 18:50:52 -0700 Subject: [PATCH 2/6] SSH-key-gen fixes. --- src/azure-cli-core/HISTORY.rst | 2 + src/azure-cli-core/azure/cli/core/keys.py | 29 ++++- .../azure/cli/core/tests/test_keys.py | 121 ++++++++++++++++++ .../cli/command_modules/vm/_validators.py | 15 +-- .../hybrid_2018_03_01/test_vm_actions.py | 15 ++- .../vm/tests/latest/test_vm_actions.py | 15 ++- .../profile_2017_03_09/test_vm_actions.py | 15 ++- 7 files changed, 186 insertions(+), 26 deletions(-) create mode 100644 src/azure-cli-core/azure/cli/core/tests/test_keys.py diff --git a/src/azure-cli-core/HISTORY.rst b/src/azure-cli-core/HISTORY.rst index b52e5ce6347..f85bfdc786b 100644 --- a/src/azure-cli-core/HISTORY.rst +++ b/src/azure-cli-core/HISTORY.rst @@ -6,6 +6,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, #6790) 2.0.45 ++++++ diff --git a/src/azure-cli-core/azure/cli/core/keys.py b/src/azure-cli-core/azure/cli/core/keys.py index 7e609e909ea..5e1eb8043a3 100644 --- a/src/azure-cli-core/azure/cli/core/keys.py +++ b/src/azure-cli-core/azure/cli/core/keys.py @@ -6,6 +6,9 @@ import os import os.path +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 @@ -32,14 +35,34 @@ def is_valid_ssh_rsa_public_key(openssh_pubkey): def generate_ssh_keys(private_key_filepath, public_key_filepath): import paramiko + try: + with open(public_key_filepath, 'r') as public_key_file: + public_key = public_key_file.read() + pub_ssh_dir, _ = os.path.split(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) + + return public_key + except IOError: + pass + ssh_dir, _ = os.path.split(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. + key = paramiko.RSAKey(filename=private_key_filepath) + logger.warning("Private SSH key file '%s' found in dir '%s'." + " Generating public key file '%s'", + private_key_filepath, ssh_dir, public_key_filepath) + 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()) diff --git a/src/azure-cli-core/azure/cli/core/tests/test_keys.py b/src/azure-cli-core/azure/cli/core/tests/test_keys.py new file mode 100644 index 00000000000..ce36fb12064 --- /dev/null +++ b/src/azure-cli-core/azure/cli/core/tests/test_keys.py @@ -0,0 +1,121 @@ +# -------------------------------------------------------------------------------------------- +# 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 azure.cli.core.keys import generate_ssh_keys + + +class TestGenerateSSHKeys(unittest.TestCase): + + @classmethod + def setUpClass(cls): + # set up temporary directory to be used for temp files. + cls._tempdirName = tempfile.mkdtemp(prefix="key_tmp_") + + cls.key = paramiko.RSAKey.generate(2048) + keyOutput = io.StringIO() + cls.key.write_private_key(keyOutput) + + cls.private_key = keyOutput.getvalue() + cls.public_key = '{} {}'.format(cls.key.get_name(), cls.key.get_base64()) + + @classmethod + def tearDownClass(cls): + # delete temporary directory to be used for temp files. + shutil.rmtree(cls._tempdirName) + + def tearDown(self): + + # remove all temporary files/directories created by previous test method. + dir = self._tempdirName + file_paths = [os.path.join(dir, name) for name in os.listdir(dir)] + + for fp in file_paths: + if os.path.isfile(fp): + os.remove(fp) + else: + shutil.rmtree(fp) + + def test_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_private_key_file_exists(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_private_key_file_new(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) + file_path = f.name + + return file_path + + +if __name__ == '__main__': + unittest.main() diff --git a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_validators.py b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_validators.py index c40f64dad61..e4160d4e34b 100644 --- a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_validators.py +++ b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/_validators.py @@ -864,28 +864,15 @@ def _validate_admin_password(password, os_type): def validate_ssh_key(namespace): - - # check if id_rsa is present but id_rsa.pub is missing - def _public_key_is_missing(): - private_key_path = os.path.join(os.path.expanduser('~'), '.ssh', 'id_rsa') - public_key_path = os.path.join(os.path.expanduser('~'), '.ssh', 'id_rsa.pub') - return (os.path.exists(private_key_path) and not os.path.exists(public_key_path)) - - - string_or_file = (namespace.ssh_key_value or os.path.join(os.path.expanduser('~'), '.ssh', 'id_rsa.pub')) content = string_or_file if os.path.exists(string_or_file): - logger.info('Reusing existing SSH public key from %s', string_or_file) + logger.info('Use existing SSH public key file: %s', string_or_file) with open(string_or_file, 'r') as f: content = f.read() elif not keys.is_valid_ssh_rsa_public_key(content): if namespace.generate_ssh_keys: - - if _public_key_is_missing(): - raise CLIError('SSH private key id_rsa exists but public key is missing. Please export the public key.') - # figure out appropriate file names: # 'base_name'(with private keys), and 'base_name.pub'(with public keys) public_key_filepath = string_or_file diff --git a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/tests/hybrid_2018_03_01/test_vm_actions.py b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/tests/hybrid_2018_03_01/test_vm_actions.py index f013e002414..eacd4c894a9 100644 --- a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/tests/hybrid_2018_03_01/test_vm_actions.py +++ b/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,13 @@ 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_") + + # 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 +57,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 +66,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 @@ -69,6 +75,9 @@ def test_generate_specfied_ssh_key_files(self): self.assertTrue(os.path.isfile(public_key_file4 + '.private')) self.assertTrue(os.path.isfile(public_key_file4)) + # delete temporary directory and its files + shutil.rmtree(temp_dir_name) + def test_figure_out_storage_source(self): test_data = 'https://av123images.blob.core.windows.net/images/TDAZBET.vhd' src_blob_uri, src_disk, src_snapshot = _figure_out_storage_source(DummyCli(), 'tg1', test_data) diff --git a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/tests/latest/test_vm_actions.py b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/tests/latest/test_vm_actions.py index f013e002414..eacd4c894a9 100644 --- a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/tests/latest/test_vm_actions.py +++ b/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,13 @@ 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_") + + # 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 +57,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 +66,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 @@ -69,6 +75,9 @@ def test_generate_specfied_ssh_key_files(self): self.assertTrue(os.path.isfile(public_key_file4 + '.private')) self.assertTrue(os.path.isfile(public_key_file4)) + # delete temporary directory and its files + shutil.rmtree(temp_dir_name) + def test_figure_out_storage_source(self): test_data = 'https://av123images.blob.core.windows.net/images/TDAZBET.vhd' src_blob_uri, src_disk, src_snapshot = _figure_out_storage_source(DummyCli(), 'tg1', test_data) diff --git a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/tests/profile_2017_03_09/test_vm_actions.py b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/tests/profile_2017_03_09/test_vm_actions.py index a74d6825607..23f6810941d 100644 --- a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/tests/profile_2017_03_09/test_vm_actions.py +++ b/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,13 @@ 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_") + + # 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 +56,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 +65,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 @@ -68,6 +74,9 @@ def test_generate_specfied_ssh_key_files(self): self.assertTrue(os.path.isfile(public_key_file4 + '.private')) self.assertTrue(os.path.isfile(public_key_file4)) + # delete temporary directory and its files + shutil.rmtree(temp_dir_name) + def test_figure_out_storage_source(self): test_data = 'https://av123images.blob.core.windows.net/images/TDAZBET.vhd' src_blob_uri, src_disk, src_snapshot = _figure_out_storage_source(DummyCli(), 'tg1', test_data) From 285a90aca12729eccb598c352c5ea9870cae9cdb Mon Sep 17 00:00:00 2001 From: Oluwatosin Adewale Date: Tue, 4 Sep 2018 14:45:05 -0700 Subject: [PATCH 3/6] Fixed Pep8 issue. --- src/azure-cli-core/azure/cli/core/tests/test_keys.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/azure-cli-core/azure/cli/core/tests/test_keys.py b/src/azure-cli-core/azure/cli/core/tests/test_keys.py index ce36fb12064..64d28fc3232 100644 --- a/src/azure-cli-core/azure/cli/core/tests/test_keys.py +++ b/src/azure-cli-core/azure/cli/core/tests/test_keys.py @@ -108,7 +108,6 @@ def test_private_key_file_new(self): 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) From 71f15f3d7f5e3cb5e590461a194746bff6e44241 Mon Sep 17 00:00:00 2001 From: Oluwatosin Adewale Date: Tue, 4 Sep 2018 18:17:39 -0700 Subject: [PATCH 4/6] Addressed PR comments. Added new test methods to TestGenerateSSHKeys. --- src/azure-cli-core/azure/cli/core/keys.py | 45 ++++++++----- .../azure/cli/core/tests/test_keys.py | 65 ++++++++++++------- .../hybrid_2018_03_01/test_vm_actions.py | 6 +- .../vm/tests/latest/test_vm_actions.py | 6 +- .../profile_2017_03_09/test_vm_actions.py | 6 +- 5 files changed, 76 insertions(+), 52 deletions(-) diff --git a/src/azure-cli-core/azure/cli/core/keys.py b/src/azure-cli-core/azure/cli/core/keys.py index 5e1eb8043a3..06063679edc 100644 --- a/src/azure-cli-core/azure/cli/core/keys.py +++ b/src/azure-cli-core/azure/cli/core/keys.py @@ -6,6 +6,7 @@ import os import os.path +from knack.util import CLIError from knack.log import get_logger logger = get_logger(__name__) @@ -34,30 +35,38 @@ 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 - try: - with open(public_key_filepath, 'r') as public_key_file: - public_key = public_key_file.read() - pub_ssh_dir, _ = os.path.split(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) - - return public_key - except IOError: - pass - - ssh_dir, _ = os.path.split(private_key_filepath) + 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) + + 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) if os.path.isfile(private_key_filepath): # try to use existing private key if it exists. - key = paramiko.RSAKey(filename=private_key_filepath) - logger.warning("Private SSH key file '%s' found in dir '%s'." - " Generating public key file '%s'", - private_key_filepath, ssh_dir, public_key_filepath) + + 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) @@ -65,7 +74,7 @@ def generate_ssh_keys(private_key_filepath, public_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) diff --git a/src/azure-cli-core/azure/cli/core/tests/test_keys.py b/src/azure-cli-core/azure/cli/core/tests/test_keys.py index 64d28fc3232..e9dbf270b54 100644 --- a/src/azure-cli-core/azure/cli/core/tests/test_keys.py +++ b/src/azure-cli-core/azure/cli/core/tests/test_keys.py @@ -10,43 +10,29 @@ import io import os import shutil +from knack.util import CLIError from azure.cli.core.keys import generate_ssh_keys class TestGenerateSSHKeys(unittest.TestCase): - @classmethod - def setUpClass(cls): + def setUp(self): # set up temporary directory to be used for temp files. - cls._tempdirName = tempfile.mkdtemp(prefix="key_tmp_") + self._tempdirName = tempfile.mkdtemp(prefix="key_tmp_") - cls.key = paramiko.RSAKey.generate(2048) + self.key = paramiko.RSAKey.generate(2048) keyOutput = io.StringIO() - cls.key.write_private_key(keyOutput) + self.key.write_private_key(keyOutput) - cls.private_key = keyOutput.getvalue() - cls.public_key = '{} {}'.format(cls.key.get_name(), cls.key.get_base64()) - - @classmethod - def tearDownClass(cls): - # delete temporary directory to be used for temp files. - shutil.rmtree(cls._tempdirName) + self.private_key = keyOutput.getvalue() + self.public_key = '{} {}'.format(self.key.get_name(), self.key.get_base64()) def tearDown(self): - - # remove all temporary files/directories created by previous test method. - dir = self._tempdirName - file_paths = [os.path.join(dir, name) for name in os.listdir(dir)] - - for fp in file_paths: - if os.path.isfile(fp): - os.remove(fp) - else: - shutil.rmtree(fp) + # delete temporary directory to be used for temp files. + shutil.rmtree(self._tempdirName) def test_public_key_file_exists(self): - # Create public key file public_key_path = self._create_new_temp_key_file(self.public_key, suffix=".pub") @@ -66,8 +52,38 @@ def test_public_key_file_exists(self): new_private_key = f.read() self.assertEqual(self.private_key, new_private_key) - def test_private_key_file_exists(self): + def test_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_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_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_private_key_file_exists(self): # Create private key file private_key_path = self._create_new_temp_key_file(self.private_key) @@ -87,7 +103,6 @@ def test_private_key_file_exists(self): self.assertEqual(self.private_key, private_key) def test_private_key_file_new(self): - # create random temp file name f = tempfile.NamedTemporaryFile(mode='w', dir=self._tempdirName) f.close() diff --git a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/tests/hybrid_2018_03_01/test_vm_actions.py b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/tests/hybrid_2018_03_01/test_vm_actions.py index eacd4c894a9..dc5192efcd4 100644 --- a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/tests/hybrid_2018_03_01/test_vm_actions.py +++ b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/tests/hybrid_2018_03_01/test_vm_actions.py @@ -30,6 +30,9 @@ class TestActions(unittest.TestCase): def test_generate_specfied_ssh_key_files(self): 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' @@ -75,9 +78,6 @@ def test_generate_specfied_ssh_key_files(self): self.assertTrue(os.path.isfile(public_key_file4 + '.private')) self.assertTrue(os.path.isfile(public_key_file4)) - # delete temporary directory and its files - shutil.rmtree(temp_dir_name) - def test_figure_out_storage_source(self): test_data = 'https://av123images.blob.core.windows.net/images/TDAZBET.vhd' src_blob_uri, src_disk, src_snapshot = _figure_out_storage_source(DummyCli(), 'tg1', test_data) diff --git a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/tests/latest/test_vm_actions.py b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/tests/latest/test_vm_actions.py index eacd4c894a9..dc5192efcd4 100644 --- a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/tests/latest/test_vm_actions.py +++ b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/tests/latest/test_vm_actions.py @@ -30,6 +30,9 @@ class TestActions(unittest.TestCase): def test_generate_specfied_ssh_key_files(self): 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' @@ -75,9 +78,6 @@ def test_generate_specfied_ssh_key_files(self): self.assertTrue(os.path.isfile(public_key_file4 + '.private')) self.assertTrue(os.path.isfile(public_key_file4)) - # delete temporary directory and its files - shutil.rmtree(temp_dir_name) - def test_figure_out_storage_source(self): test_data = 'https://av123images.blob.core.windows.net/images/TDAZBET.vhd' src_blob_uri, src_disk, src_snapshot = _figure_out_storage_source(DummyCli(), 'tg1', test_data) diff --git a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/tests/profile_2017_03_09/test_vm_actions.py b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/tests/profile_2017_03_09/test_vm_actions.py index 23f6810941d..9d02253d1ac 100644 --- a/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/tests/profile_2017_03_09/test_vm_actions.py +++ b/src/command_modules/azure-cli-vm/azure/cli/command_modules/vm/tests/profile_2017_03_09/test_vm_actions.py @@ -29,6 +29,9 @@ class TestActions(unittest.TestCase): def test_generate_specfied_ssh_key_files(self): 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' @@ -74,9 +77,6 @@ def test_generate_specfied_ssh_key_files(self): self.assertTrue(os.path.isfile(public_key_file4 + '.private')) self.assertTrue(os.path.isfile(public_key_file4)) - # delete temporary directory and its files - shutil.rmtree(temp_dir_name) - def test_figure_out_storage_source(self): test_data = 'https://av123images.blob.core.windows.net/images/TDAZBET.vhd' src_blob_uri, src_disk, src_snapshot = _figure_out_storage_source(DummyCli(), 'tg1', test_data) From 4983dc56db7e83cfdc7bd7f2444d25eb3573ada5 Mon Sep 17 00:00:00 2001 From: Oluwatosin Adewale Date: Wed, 5 Sep 2018 11:42:21 -0700 Subject: [PATCH 5/6] Addressed more comments. --- src/azure-cli-core/HISTORY.rst | 2 +- .../azure/cli/core/tests/test_keys.py | 16 +++++++--------- 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/src/azure-cli-core/HISTORY.rst b/src/azure-cli-core/HISTORY.rst index f85bfdc786b..8816317ea8a 100644 --- a/src/azure-cli-core/HISTORY.rst +++ b/src/azure-cli-core/HISTORY.rst @@ -7,7 +7,7 @@ Release History ++++++ * Minor fixes. * Fixed issue where `az vm create --generate-ssh-keys` overwrites private key - file if public key file is missing. (#4725, #6790) + file if public key file is missing. (#4725, #6780) 2.0.45 ++++++ diff --git a/src/azure-cli-core/azure/cli/core/tests/test_keys.py b/src/azure-cli-core/azure/cli/core/tests/test_keys.py index e9dbf270b54..9083e4539d7 100644 --- a/src/azure-cli-core/azure/cli/core/tests/test_keys.py +++ b/src/azure-cli-core/azure/cli/core/tests/test_keys.py @@ -32,7 +32,7 @@ def tearDown(self): # delete temporary directory to be used for temp files. shutil.rmtree(self._tempdirName) - def test_public_key_file_exists(self): + 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") @@ -52,7 +52,7 @@ def test_public_key_file_exists(self): new_private_key = f.read() self.assertEqual(self.private_key, new_private_key) - def test_public_key_file_exists_no_permissions(self): + 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) @@ -61,7 +61,7 @@ def test_public_key_file_exists_no_permissions(self): with self.assertRaises(CLIError): generate_ssh_keys("", public_key_path) - def test_private_key_file_exists_no_permissions(self): + 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) @@ -71,7 +71,7 @@ def test_private_key_file_exists_no_permissions(self): public_key_path = private_key_path + ".pub" generate_ssh_keys(private_key_path, public_key_path) - def test_private_key_file_exists_encrypted(self): + 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("") @@ -83,7 +83,7 @@ def test_private_key_file_exists_encrypted(self): public_key_path = private_key_path + ".pub" generate_ssh_keys(private_key_path, public_key_path) - def test_private_key_file_exists(self): + 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) @@ -102,7 +102,7 @@ def test_private_key_file_exists(self): private_key = f.read() self.assertEqual(self.private_key, private_key) - def test_private_key_file_new(self): + def test_generate_new_private_public_key_files(self): # create random temp file name f = tempfile.NamedTemporaryFile(mode='w', dir=self._tempdirName) f.close() @@ -126,9 +126,7 @@ def test_private_key_file_new(self): 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) - file_path = f.name - - return file_path + return f.name if __name__ == '__main__': From 0cb24eb057b0ffb660abc3d3c16ab39b1411f4a4 Mon Sep 17 00:00:00 2001 From: Oluwatosin Adewale Date: Thu, 6 Sep 2018 14:01:36 -0700 Subject: [PATCH 6/6] Addressed some more comments. --- src/azure-cli-core/HISTORY.rst | 1 - src/azure-cli-core/azure/cli/core/keys.py | 2 -- 2 files changed, 3 deletions(-) diff --git a/src/azure-cli-core/HISTORY.rst b/src/azure-cli-core/HISTORY.rst index 8816317ea8a..8816af0d7e4 100644 --- a/src/azure-cli-core/HISTORY.rst +++ b/src/azure-cli-core/HISTORY.rst @@ -5,7 +5,6 @@ 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) diff --git a/src/azure-cli-core/azure/cli/core/keys.py b/src/azure-cli-core/azure/cli/core/keys.py index 06063679edc..89cb9ec9adc 100644 --- a/src/azure-cli-core/azure/cli/core/keys.py +++ b/src/azure-cli-core/azure/cli/core/keys.py @@ -57,13 +57,11 @@ def generate_ssh_keys(private_key_filepath, public_key_filepath): 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)