def _CreateSshKeyPairIfNecessary(cfg):
"""Create ssh key pair if necessary.
Args:
cfg: An Acloudconfig instance.
Raises:
error.DriverError: If it falls into an unexpected condition.
"""
if not cfg.ssh_public_key_path:
logger.warning("ssh_public_key_path is not specified in acloud config. "
"Project-wide public key will "
"be used when creating AVD instances. "
"Please ensure you have the correct private half of "
"a project-wide public key if you want to ssh into the "
"instances after creation.")
elif cfg.ssh_public_key_path and not cfg.ssh_private_key_path:
logger.warning("Only ssh_public_key_path is specified in acloud config,"
" but ssh_private_key_path is missing. "
"Please ensure you have the correct private half "
"if you want to ssh into the instances after creation.")
elif cfg.ssh_public_key_path and cfg.ssh_private_key_path:
utils.CreateSshKeyPairIfNotExist(
cfg.ssh_private_key_path, cfg.ssh_public_key_path)
else:
# Should never reach here.
raise errors.DriverError(
"Unexpected error in _CreateSshKeyPairIfNecessary")
def testCreateSshKeyPairKeyAlreadyExists(self): #pylint: disable=invalid-name
"""Test when the key pair already exists."""
public_key = "/fake/public_key"
private_key = "/fake/private_key"
self.Patch(os.path, "exists", side_effect=[True, True])
self.Patch(subprocess, "check_call")
self.Patch(os, "makedirs", return_value=True)
utils.CreateSshKeyPairIfNotExist(private_key, public_key)
self.assertEqual(subprocess.check_call.call_count, 0) #pylint: disable=no-member
def testCreateSshKeyPair_KeyAlreadyExists(self):
"""Test when the key pair already exists."""
public_key = "/fake/public_key"
private_key = "/fake/private_key"
self.Patch(os.path,
"exists",
side_effect=lambda path: path == public_key)
self.Patch(subprocess, "check_call")
utils.CreateSshKeyPairIfNotExist(private_key, public_key)
self.assertEqual(subprocess.check_call.call_count, 0)
def testCreateSshKeyPair_KeyAreCreated(self):
"""Test when the key pair created."""
public_key = "/fake/public_key"
private_key = "/fake/private_key"
self.Patch(os.path, "exists", return_value=False)
self.Patch(subprocess, "check_call")
self.Patch(os, "rename")
utils.CreateSshKeyPairIfNotExist(private_key, public_key)
self.assertEqual(subprocess.check_call.call_count, 1)
subprocess.check_call.assert_called_with(
utils.SSH_KEYGEN_CMD +
["-C", getpass.getuser(), "-f", private_key],
stdout=mock.ANY,
stderr=mock.ANY)
def testCreatePublicKeyAreCreated(self):
"""Test when the PublicKey created."""
public_key = "/fake/public_key"
private_key = "/fake/private_key"
self.Patch(os.path, "exists", side_effect=[False, True, True])
self.Patch(os, "makedirs", return_value=True)
mock_open = mock.mock_open(read_data=public_key)
self.Patch(subprocess, "check_output")
self.Patch(os, "rename")
with mock.patch.object(six.moves.builtins, "open", mock_open):
utils.CreateSshKeyPairIfNotExist(private_key, public_key)
self.assertEqual(subprocess.check_output.call_count, 1) #pylint: disable=no-member
subprocess.check_output.assert_called_with( #pylint: disable=no-member
utils.SSH_KEYGEN_PUB_CMD + ["-f", private_key])
def SetupSSHKeys(config_path, private_key_path, public_key_path):
"""Setup the pair of the ssh key for acloud.config.
User can use the default path: "~/.ssh/acloud_rsa".
Args:
config_path: String, acloud config path.
private_key_path: Path to the private key file.
e.g. ~/.ssh/acloud_rsa
public_key_path: Path to the public key file.
e.g. ~/.ssh/acloud_rsa.pub
"""
private_key_path = os.path.expanduser(private_key_path)
if (private_key_path == "" or public_key_path == ""
or private_key_path == _DEFAULT_SSH_PRIVATE_KEY):
utils.CreateSshKeyPairIfNotExist(_DEFAULT_SSH_PRIVATE_KEY,
_DEFAULT_SSH_PUBLIC_KEY)
UpdateConfigFile(config_path, "ssh_private_key_path",
_DEFAULT_SSH_PRIVATE_KEY)
UpdateConfigFile(config_path, "ssh_public_key_path",
_DEFAULT_SSH_PUBLIC_KEY)