Ssh cleanup

src/ssh/HISTORY.md

@@ -1,5 +1,14 @@
 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.
+* Add support to connect to local users on local machines using key based, cert based, or password based authentication.
+
 0.1.8
 -----
 * Rollback from version 0.1.7 to 0.1.6 to remove preview features.

src/ssh/azext_ssh/_help.py

@@ -13,20 +13,36 @@
 helps['ssh vm'] = """
     type: command
     short-summary: SSH into Azure VMs using an ssh certificate
+    long-summary: Users can login using AAD issued certificates or using local user credentials. We recommend login using AAD issued certificates. To SSH as a local user in the target machine, you must provide the local user name using the --local-user argument.
     examples:
         - name: Give a resource group and VM to SSH to
           text: |
             az ssh vm --resource-group myResourceGroup --vm-name myVm
+
         - name: Give the public IP (or hostname) of a VM to SSH to
           text: |
             az ssh vm --ip 1.2.3.4
             az ssh vm --hostname example.com
+
         - name: Using a custom private key file
           text: |
             az ssh vm --ip 1.2.3.4 --private-key-file key --public-key-file key.pub
+
         - name: Using additional ssh arguments
           text: |
             az ssh vm --ip 1.2.3.4 -- -A -o ForwardX11=yes
+
+        - name: Give a local user name to SSH using local user credentials on the target machine using certificate based authentication.
+          text: |
+            az ssh vm --local-user username --ip 1.2.3.4 --certificate-file cert.pub --private-key key
+
+        - name: Give a local user name to SSH using local user credentials on the target machine using key based authentication.
+          text: |
+            az ssh vm --local-user username --resource-group myResourceGroup --vm-name myVM --private-key-file key
+
+        - name: Give a local user name to SSH using local user credentials on the target machine using password based authentication.
+          text: |
+            az ssh vm --local-user username --ip 1.2.3.4
 """
 
 helps['ssh config'] = """
@@ -37,16 +53,19 @@
         - name: Give a resource group and VM for which to create a config, and save in a local file
           text: |
             az ssh config --resource-group myResourceGroup --vm-name myVm --file ./sshconfig
+
         - name: Give the public IP (or hostname) of a VM for which to create a config and then ssh
           text: |
             az ssh config --ip 1.2.3.4 --file ./sshconfig
             ssh -F ./sshconfig 1.2.3.4
+
         - name: Create a generic config for use with any host
           text: |
             #Bash
             az ssh config --ip \\* --file ./sshconfig
             #PowerShell
             az ssh config --ip * --file ./sshconfig
+
         - name: Examples with other software
           text: |
             #Bash

src/ssh/azext_ssh/_params.py

@@ -14,6 +14,10 @@ def load_arguments(self, _):
         c.argument('private_key_file', options_list=['--private-key-file', '-i'], help='The RSA private key file path')
         c.argument('use_private_ip', options_list=['--prefer-private-ip'],
                    help='Will prefer private IP.  Requires connectivity to the private IP.')
+        c.argument('local_user', options_list=['--local-user'],
+                   help='The username for a local user')
+        c.argument('cert_file', options_list=['--certificate-file', '-c'],
+                   help='Path to a certificate file used for authentication when using local user credentials.')
         c.argument('port', options_list=['--port'], help='SSH port')
         c.positional('ssh_args', nargs='*', help='Additional arguments passed to OpenSSH')
 
@@ -25,10 +29,18 @@ def load_arguments(self, _):
         c.argument('private_key_file', options_list=['--private-key-file', '-i'], help='The RSA private key file path')
         c.argument('use_private_ip', options_list=['--prefer-private-ip'],
                    help='Will use a private IP if available. By default only public IPs are used.')
+        c.argument('local_user', options_list=['--local-user'],
+                   help='The username for a local user')
         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', '--keys-dest-folder'],
+                   help='Folder where new generated keys will be stored.')
+        c.argument('cert_file', options_list=['--certificate-file', '-c'], help='Path to certificate file')
+        c.argument('port', options_list=['--port'], help='SSH port')
 
     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):
+           private_key_file=None, use_private_ip=False, local_user=None, cert_file=None, port=None,
+           ssh_args=None):
+    _assert_args(resource_group_name, vm_name, ssh_ip, cert_file, local_user)
+    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)
+    _do_ssh_op(cmd, resource_group_name, vm_name, ssh_ip, public_key_file, private_key_file, use_private_ip,
+               local_user, cert_file, 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):
-    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)
+               public_key_file=None, private_key_file=None, overwrite=False, use_private_ip=False,
+               local_user=None, cert_file=None, port=None, credentials_folder=None):
+    _assert_args(resource_group_name, vm_name, ssh_ip, cert_file, local_user)
+    # 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, port)
+    # Default credential location
+    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,
+               local_user, cert_file, 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,
+               username, cert_file, 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,18 @@ 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")
 
-    cert_file, username = _get_and_write_certificate(cmd, public_key_file, None)
-    op_call(ssh_ip, username, cert_file, private_key_file)
+    # If user provides local user, no credentials should be deleted.
+    delete_keys = False
+    delete_cert = False
+    # If user provides a local user, use the provided credentials for authentication
+    if not username:
+        delete_cert = True
+        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, delete_keys, delete_cert)
 
 
 def _get_and_write_certificate(cmd, public_key_file, cert_file):
@@ -106,7 +151,7 @@ def _prepare_jwk_data(public_key_file):
     return data
 
 
-def _assert_args(resource_group, vm_name, ssh_ip):
+def _assert_args(resource_group, vm_name, ssh_ip, cert_file, username):
     if not (resource_group or vm_name or ssh_ip):
         raise azclierror.RequiredArgumentMissingError(
             "The VM must be specified by --ip or --resource-group and --vm-name/--name")
@@ -119,13 +164,31 @@ def _assert_args(resource_group, vm_name, ssh_ip):
         raise azclierror.MutuallyExclusiveArgumentError(
             "--ip cannot be used with --resource-group or --vm-name/--name")
 
+    if cert_file and not username:
+        raise azclierror.MutuallyExclusiveArgumentError(
+            "To authenticate with a certificate you need to provide a --local-user")
 
-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
+    if cert_file and not os.path.isfile(cert_file):
+        raise azclierror.FileOperationError(f"Certificate file {cert_file} not found")
+
+
+def _check_or_create_public_private_files(public_key_file, private_key_file, credentials_folder):
+    delete_keys = False
+    # If nothing is passed, then 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 provided 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 +206,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