Ssh cleanup

src/ssh/HISTORY.md

@@ -1,5 +1,13 @@
 Release History
 ===============
+1.0.0
+-----
+* Delete all keys and certificates created during execution of ssh vm.
+* Add --keys-destination-folder to ssh config
+* Keys generated during ssh config are saved in az_ssh_config folder in the same directory as --file.
+* Users no longer allowed to run ssh cert with no parameters. 
+* When --public-key-file/-f is not provided to ssh cert, generated public and private keys are saved in the same folder as --file.
+
 0.1.8
 -----
 * Rollback from version 0.1.7 to 0.1.6 to remove preview features.

src/ssh/azext_ssh/_params.py

@@ -27,8 +27,12 @@ def load_arguments(self, _):
                    help='Will use a private IP if available. By default only public IPs are used.')
         c.argument('overwrite', action='store_true', options_list=['--overwrite'],
                    help='Overwrites the config file if this flag is set')
+        c.argument('credentials_folder', options_list=['--keys-destination-folder'],
+                   help='Folder where new generated keys will be stored.')
 
     with self.argument_context('ssh cert') as c:
         c.argument('cert_path', options_list=['--file', '-f'],
                    help='The file path to write the SSH cert to, defaults to public key path with -aadcert.pub appened')
-        c.argument('public_key_file', options_list=['--public-key-file', '-p'], help='The RSA public key file path')
+        c.argument('public_key_file', options_list=['--public-key-file', '-p'],
+                   help='The RSA public key file path. If not provided, '
+                   'generated key pair is stored in the same directory as --file')

src/ssh/azext_ssh/custom.py

@@ -9,35 +9,70 @@
 import json
 import tempfile
 
+from knack import log
 from azure.cli.core import azclierror
 
 from . import ip_utils
 from . import rsa_parser
 from . import ssh_utils
 
+logger = log.get_logger(__name__)
+
 
 def ssh_vm(cmd, resource_group_name=None, vm_name=None, ssh_ip=None, public_key_file=None,
            private_key_file=None, use_private_ip=False, port=None, ssh_args=None):
+    _assert_args(resource_group_name, vm_name, ssh_ip)
+    credentials_folder = None
     op_call = functools.partial(ssh_utils.start_ssh_connection, port, ssh_args)
     _do_ssh_op(cmd, resource_group_name, vm_name, ssh_ip,
-               public_key_file, private_key_file, use_private_ip, op_call)
+               public_key_file, private_key_file, use_private_ip, credentials_folder, op_call)
 
 
 def ssh_config(cmd, config_path, resource_group_name=None, vm_name=None, ssh_ip=None,
-               public_key_file=None, private_key_file=None, overwrite=False, use_private_ip=False):
+               public_key_file=None, private_key_file=None, overwrite=False, use_private_ip=False,
+               credentials_folder=None):
+    _assert_args(resource_group_name, vm_name, ssh_ip)
+    # If user provides their own key pair, certificate will be written in the same folder as public key.
+    if (public_key_file or private_key_file) and credentials_folder:
+        raise azclierror.ArgumentUsageError("--keys-destination-folder can't be used in conjunction with "
+                                            "--public-key-file/-p or --private-key-file/-i.")
     op_call = functools.partial(ssh_utils.write_ssh_config, config_path, resource_group_name, vm_name, overwrite)
-    _do_ssh_op(cmd, resource_group_name, vm_name, ssh_ip, public_key_file, private_key_file, use_private_ip, op_call)
+    # Default credential location
+    if credentials_folder and not os.path.isdir(credentials_folder):
+        raise azclierror.InvalidArgumentValueError(f"--keys-destination-folder {credentials_folder} doesn't exist")
+    if not credentials_folder:
+        config_folder = os.path.dirname(config_path)
+        if not os.path.isdir(config_folder):
+            raise azclierror.InvalidArgumentValueError(f"Config file destination folder {config_folder} "
+                                                       "does not exist.")
+        folder_name = ssh_ip
+        if resource_group_name and vm_name:
+            folder_name = resource_group_name + "-" + vm_name
+        credentials_folder = os.path.join(config_folder, os.path.join("az_ssh_config", folder_name))
+
+    _do_ssh_op(cmd, resource_group_name, vm_name, ssh_ip, public_key_file, private_key_file, use_private_ip,
+               credentials_folder, op_call)
 
 
 def ssh_cert(cmd, cert_path=None, public_key_file=None):
-    public_key_file, _ = _check_or_create_public_private_files(public_key_file, None)
+    if not cert_path and not public_key_file:
+        raise azclierror.RequiredArgumentMissingError("--file or --public-key-file must be provided.")
+    if cert_path and not os.path.isdir(os.path.dirname(cert_path)):
+        raise azclierror.InvalidArgumentValueError(f"{os.path.dirname(cert_path)} folder doesn't exist")
+    # If user doesn't provide a public key, save generated key pair to the same folder as --file
+    keys_folder = None
+    if not public_key_file:
+        keys_folder = os.path.dirname(cert_path)
+        logger.warning("The generated SSH keys are stored at %s. Please delete SSH keys when the certificate "
+                       "is no longer being used.", keys_folder)
+    public_key_file, _, _ = _check_or_create_public_private_files(public_key_file, None, keys_folder)
     cert_file, _ = _get_and_write_certificate(cmd, public_key_file, cert_path)
     print(cert_file + "\n")
 
 
-def _do_ssh_op(cmd, resource_group, vm_name, ssh_ip, public_key_file, private_key_file, use_private_ip, op_call):
-    _assert_args(resource_group, vm_name, ssh_ip)
-    public_key_file, private_key_file = _check_or_create_public_private_files(public_key_file, private_key_file)
+def _do_ssh_op(cmd, resource_group, vm_name, ssh_ip, public_key_file, private_key_file, use_private_ip,
+               credentials_folder, op_call):
+    # Get ssh_ip before getting public key to avoid getting "ResourceNotFound" exception after creating the keys
     ssh_ip = ssh_ip or ip_utils.get_ssh_ip(cmd, resource_group, vm_name, use_private_ip)
 
     if not ssh_ip:
@@ -46,8 +81,11 @@ def _do_ssh_op(cmd, resource_group, vm_name, ssh_ip, public_key_file, private_ke
 
         raise azclierror.ResourceNotFoundError(f"VM '{vm_name}' does not have a public or private IP address to SSH to")
 
+    public_key_file, private_key_file, delete_keys = _check_or_create_public_private_files(public_key_file,
+                                                                                           private_key_file,
+                                                                                           credentials_folder)
     cert_file, username = _get_and_write_certificate(cmd, public_key_file, None)
-    op_call(ssh_ip, username, cert_file, private_key_file)
+    op_call(ssh_ip, username, cert_file, private_key_file, delete_keys)
 
 
 def _get_and_write_certificate(cmd, public_key_file, cert_file):
@@ -120,12 +158,23 @@ def _assert_args(resource_group, vm_name, ssh_ip):
             "--ip cannot be used with --resource-group or --vm-name/--name")
 
 
-def _check_or_create_public_private_files(public_key_file, private_key_file):
-    # If nothing is passed in create a temporary directory with a ephemeral keypair
+def _check_or_create_public_private_files(public_key_file, private_key_file, credentials_folder):
+    delete_keys = False
+    # If nothing is passed in create a directory with a ephemeral keypair
     if not public_key_file and not private_key_file:
-        temp_dir = tempfile.mkdtemp(prefix="aadsshcert")
-        public_key_file = os.path.join(temp_dir, "id_rsa.pub")
-        private_key_file = os.path.join(temp_dir, "id_rsa")
+        # We only want to delete the keys if the user hasn't providede their own keys
+        # Only ssh vm deletes generated keys.
+        delete_keys = True
+        if not credentials_folder:
+            # az ssh vm: Create keys on temp folder and delete folder once connection succeeds/fails.
+            credentials_folder = tempfile.mkdtemp(prefix="aadsshcert")
+        else:
+            # az ssh config: Keys saved to the same folder as --file or to --keys-destination-folder.
+            # az ssh cert: Keys saved to the same folder as --file.
+            if not os.path.isdir(credentials_folder):
+                os.makedirs(credentials_folder)
+        public_key_file = os.path.join(credentials_folder, "id_rsa.pub")
+        private_key_file = os.path.join(credentials_folder, "id_rsa")
         ssh_utils.create_ssh_keyfile(private_key_file)
 
     if not public_key_file:
@@ -143,7 +192,7 @@ def _check_or_create_public_private_files(public_key_file, private_key_file):
         if not os.path.isfile(private_key_file):
             raise azclierror.FileOperationError(f"Private key file {private_key_file} not found")
 
-    return public_key_file, private_key_file
+    return public_key_file, private_key_file, delete_keys
 
 
 def _write_cert_file(certificate_contents, cert_file):

src/ssh/azext_ssh/file_utils.py

@@ -5,6 +5,11 @@
 
 import errno
 import os
+from knack import log
+
+from azure.cli.core import azclierror
+
+logger = log.get_logger(__name__)
 
 
 def make_dirs_for_file(file_path):
@@ -20,3 +25,25 @@ def mkdir_p(path):
             pass
         else:
             raise
+
+
+def delete_file(file_path, message, warning=False):
+    if os.path.isfile(file_path):
+        try:
+            os.remove(file_path)
+        except Exception as e:
+            if warning:
+                logger.warning(message)
+            else:
+                raise azclierror.FileOperationError(message + "Error: " + str(e)) from e
+
+
+def delete_folder(dir_path, message, warning=False):
+    if os.path.isdir(dir_path):
+        try:
+            os.rmdir(dir_path)
+        except Exception as e:
+            if warning:
+                logger.warning(message)
+            else:
+                raise azclierror.FileOperationError(message + "Error: " + str(e)) from e

src/ssh/azext_ssh/ssh_utils.py

@@ -5,22 +5,62 @@
 import os
 import platform
 import subprocess
+import time
+import multiprocessing as mp
+from azext_ssh import file_utils
 
 from knack import log
 from azure.cli.core import azclierror
 
 logger = log.get_logger(__name__)
 
+CLEANUP_TOTAL_TIME_LIMIT_IN_SECONDS = 120
+CLEANUP_TIME_INTERVAL_IN_SECONDS = 10
+CLEANUP_AWAIT_TERMINATION_IN_SECONDS = 30
+
+
+def start_ssh_connection(port, ssh_args, ip, username, cert_file, private_key_file, delete_keys):
 
-def start_ssh_connection(port, ssh_args, ip, username, cert_file, private_key_file):
     ssh_arg_list = []
     if ssh_args:
         ssh_arg_list = ssh_args
+
+    log_file = None
+    if '-E' not in ssh_arg_list and set(['-v', '-vv', '-vvv']).isdisjoint(ssh_arg_list):
+        # If the user either provides his own client log file (-E) or
+        # wants the client log messages to be printed to the console (-vvv/-vv/-v),
+        # we should not use the log files to check for connection success.
+        log_file_dir = os.path.dirname(cert_file)
+        log_file_name = 'ssh_client_log_' + str(os.getpid())
+        log_file = os.path.join(log_file_dir, log_file_name)
+        ssh_arg_list = ssh_arg_list + ['-E', log_file, '-v']
+
     command = [_get_ssh_path(), _get_host(username, ip)]
     command = command + _build_args(cert_file, private_key_file, port) + ssh_arg_list
+
+    # Create a new process that will wait until the connection is established and then delete keys.
+    cleanup_process = mp.Process(target=_do_cleanup, args=(delete_keys, cert_file, private_key_file,
+                                 log_file, True))
+    cleanup_process.start()
+
     logger.debug("Running ssh command %s", ' '.join(command))
     subprocess.call(command, shell=platform.system() == 'Windows')
 
+    if cleanup_process.is_alive():
+        cleanup_process.terminate()
+        # wait for process to terminate
+        t0 = time.time()
+        while cleanup_process.is_alive() and (time.time() - t0) < CLEANUP_AWAIT_TERMINATION_IN_SECONDS:
+            time.sleep(1)
+
+    # Make sure all files have been properly removed.
+    _do_cleanup(delete_keys, cert_file, private_key_file)
+    if log_file:
+        file_utils.delete_file(log_file, f"Couldn't delete temporary log file {cert_file}. ", True)
+    if delete_keys:
+        temp_dir = os.path.dirname(cert_file)
+        file_utils.delete_folder(temp_dir, f"Couldn't delete temporary folder {temp_dir}", True)
+
 
 def create_ssh_keyfile(private_key_file):
     command = [_get_ssh_path("ssh-keygen"), "-f", private_key_file, "-t", "rsa", "-q", "-N", ""]
@@ -50,8 +90,30 @@ def get_ssh_cert_principals(cert_file):
     return principals
 
 
+def get_ssh_cert_validity(cert_file):
+    info = get_ssh_cert_info(cert_file)
+    for line in info:
+        if "Valid:" in line:
+            return line.strip()
+    return None
+
+
 def write_ssh_config(config_path, resource_group, vm_name, overwrite,
-                     ip, username, cert_file, private_key_file):
+                     ip, username, cert_file, private_key_file, delete_keys):
+
+    # Warn users to delete credentials once config file is no longer being used.
+    # If user provided keys, only ask them to delete the certificate.
+    path_to_delete = os.path.dirname(cert_file)
+    items_to_delete = " (id_rsa, id_rsa.pub, id_rsa.pub-aadcert.pub)"
+    if not delete_keys:
+        path_to_delete = cert_file
+        items_to_delete = ""
+    validity = get_ssh_cert_validity(cert_file)
+    validity_warning = ""
+    if validity:
+        validity_warning = f" {validity.lower()}"
+    logger.warning("%s contains sensitive information%s%s, please delete it once you no longer need this config file. ",
+                   path_to_delete, items_to_delete, validity_warning)
 
     lines = [""]
 
@@ -117,3 +179,31 @@ def _build_args(cert_file, private_key_file, port):
         port_arg = ["-p", port]
     certificate = ["-o", "CertificateFile=" + cert_file]
     return private_key + certificate + port_arg
+
+
+def _do_cleanup(delete_keys, cert_file, private_key, log_file=None, wait=False):
+    # if there is a log file, use it to check for the connection success
+    if log_file:
+        t0 = time.time()
+        match = False
+        while (time.time() - t0) < CLEANUP_TOTAL_TIME_LIMIT_IN_SECONDS and not match:
+            time.sleep(CLEANUP_TIME_INTERVAL_IN_SECONDS)
+            try:
+                with open(log_file, 'r') as ssh_client_log:
+                    match = "debug1: Authentication succeeded" in ssh_client_log.read()
+                    ssh_client_log.close()
+            except:
+                # If there is an exception, wait for a little bit and try again
+                time.sleep(CLEANUP_TIME_INTERVAL_IN_SECONDS)
+
+    elif wait:
+        # if we are not checking the logs, but still want to wait for connection before deleting files
+        time.sleep(CLEANUP_TOTAL_TIME_LIMIT_IN_SECONDS)
+
+    # TO DO: Once arc changes are merged, delete relay information as well
+    if delete_keys and private_key:
+        public_key = private_key + '.pub'
+        file_utils.delete_file(private_key, f"Couldn't delete private key {private_key}. ", True)
+        file_utils.delete_file(public_key, f"Couldn't delete public key {public_key}. ", True)
+    if cert_file:
+        file_utils.delete_file(cert_file, f"Couldn't delete certificate {cert_file}. ", True)