diff --git a/src/azure-cli-core/HISTORY.rst b/src/azure-cli-core/HISTORY.rst index b52e5ce6347..8816af0d7e4 100644 --- a/src/azure-cli-core/HISTORY.rst +++ b/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 ++++++ diff --git a/src/azure-cli-core/azure/cli/core/keys.py b/src/azure-cli-core/azure/cli/core/keys.py index 7e609e909ea..89cb9ec9adc 100644 --- a/src/azure-cli-core/azure/cli/core/keys.py +++ b/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) 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..9083e4539d7 --- /dev/null +++ b/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() 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..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 @@ -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 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..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 @@ -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 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..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 @@ -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