Feature: Setup user with ssh public key instead of password (#32)

* Login with ssh public key

* Update to have ssh-key command

* Fix readme

* Fix Cr

* remove unused files

* Change head to master

README.md

@@ -3,7 +3,7 @@ A suite of distributed tools to help engineers scale their work into Azure.
 
 # Spark on DTDE
 
-## Setup  
+## Setup
 1. Clone the repo
 2. Use pip to install required packages:
 ```bash
@@ -14,9 +14,9 @@ A suite of distributed tools to help engineers scale their work into Azure.
     pip3 install .
 
     # For development use this instead
-    pip3 install -e . 
+    pip3 install -e .
 ```
-4. Rename 'configuration.cfg.template' to 'configuration.cfg' and fill in the fields for your Batch account and Storage account. These fields can be found in the Azure portal. 
+4. Rename 'configuration.cfg.template' to 'configuration.cfg' and fill in the fields for your Batch account and Storage account. These fields can be found in the Azure portal.
 
    To complete this step, you will need an Azure account that has a Batch account and Storage account:
     - To create an Azure account: https://azure.microsoft.com/free/
@@ -43,18 +43,37 @@ You can also create your cluster with [low-priority](https://docs.microsoft.com/
 ```
 azb spark cluster create \
     --id <my-cluster-id> \
-    --size-low-pri <number of low-pri nodes>
+    --size-low-pri <number of low-pri nodes> \
     --vm-size <vm-size>
 ```
 
+You can also add a user directly in this command using the same inputs as the `add-user` command described bellow.
+
+#### Add a user to your cluster to connect
 When your cluster is ready, create a user for your cluster (if you didn't already do so when creating your cluster):
-```
+```bash
+# **Recommended usage**
+# Add a user with a ssh public key. It will use the value specified in the configuration.cfg (Either path to the file or the actual key)
+azb spark cluster add-user \
+    --id <my-cluster-id> \
+    --username <username>
+
+# You can also explicity specify the ssh public key(Path or actual key)
+azb spark cluster add-user \
+    --id <my-cluster-id> \
+    --username <username> \
+    --ssh-key ~/.ssh/id_rsa.pub
+
+# **Not recommended**
+# You can also just specify a password
 azb spark cluster add-user \
     --id <my-cluster-id> \
     --username <username> \
     --password <password>
+
 ```
-NOTE: The cluster id (--id) can only contain alphanumeric characters including hyphens and underscores, and cannot contain more than 64 characters. 
+
+NOTE: The cluster id (--id) can only contain alphanumeric characters including hyphens and underscores, and cannot contain more than 64 characters.
 
 ### Submit a Spark job
 
@@ -63,8 +82,8 @@ Now you can submit jobs to run against the cluster:
 azb spark submit \
     --id <my-cluster-id> \
     --name <my-job-name> \
-    [options] 
-    <app jar | python file> 
+    [options] \
+    <app jar | python file> \
     [app arguments]
 ```
 NOTE: The job name (--name) must be atleast 3 characters long, can only contain alphanumeric characters including hyphens but excluding underscores, and cannot contain uppercase letters.
@@ -73,7 +92,7 @@ NOTE: The job name (--name) must be atleast 3 characters long, can only contain
 
 To view the spark UI, open up an ssh tunnel with the "masterui" option and a local port to map to:
 ```
-azb spark cluster ssh \ 
+azb spark cluster ssh \
     --id <my-cluster-id> \
     --masterui <local-port> \
     --username <user-name>
@@ -81,7 +100,7 @@ azb spark cluster ssh \
 
 Optionally, you can also open up a jupyter notebook with the "jupyter" option to work in:
 ```
-azb spark cluster ssh  \ 
+azb spark cluster ssh  \
     --id <my-cluster-id> \
     --masterui <local-port> \
     --jupyter <local-port>

configuration.cfg.template

@@ -9,3 +9,6 @@ storageaccountkey=
 storageaccountsuffix=core.windows.net
 
 [Default]
+# SSH public key used to create a user and connect to a server.
+# It can either be the public key itself(ssh-rsa ...) or the path to the ssh key.
+#ssh_pub_key=~/.ssh/id_rsa.pub

dtde/clusterlib.py

@@ -1,11 +1,12 @@
 from datetime import datetime, timedelta
 from subprocess import call
+from dtde.core import CommandBuilder, ssh
+from collections import namedtuple
 import azure.batch.models as batch_models
-from dtde.core import CommandBuilder
 from dtde.models import Software
+from dtde.error import ClusterNotReadyError, InvalidUserCredentialsError
 from . import azure_api, constants, upload_node_scripts, util
 
-
 POOL_ADMIN_USER_IDENTITY = batch_models.UserIdentity(
     auto_user=batch_models.AutoUserSpecification(
         scope=batch_models.AutoUserScope.pool,
@@ -24,7 +25,8 @@ def cluster_install_cmd(zip_resource_file: batch_models.ResourceFile, custom_scr
         'chmod -R 777 /usr/local/share/jupyter/kernels',
         # To avoid error: "sudo: sorry, you must have a tty to run sudo"
         'sed -i -e "s/Defaults    requiretty.*/ #Defaults    requiretty/g" /etc/sudoers',
-        'unzip $AZ_BATCH_TASK_WORKING_DIR/{0}'.format(zip_resource_file.file_path),
+        'unzip $AZ_BATCH_TASK_WORKING_DIR/{0}'.format(
+            zip_resource_file.file_path),
         'chmod 777 $AZ_BATCH_TASK_WORKING_DIR/main.sh',
         # Convert windows line ending to unix if applicable
         'dos2unix $AZ_BATCH_TASK_WORKING_DIR/main.sh',
@@ -89,7 +91,8 @@ def create_cluster(
         vm_low_pri_count,
         vm_size,
         username: str,
-        password: str,
+        password: str = None,
+        ssh_key: str = None,
         wait=True):
     """
         Create a spark cluster
@@ -136,7 +139,8 @@ def create_cluster(
         enable_inter_node_communication=True,
         max_tasks_per_node=1,
         metadata=[
-            batch_models.MetadataItem(name=constants.AZB_SOFTWARE_METADATA_KEY, value=Software.spark),
+            batch_models.MetadataItem(
+                name=constants.AZB_SOFTWARE_METADATA_KEY, value=Software.spark),
         ])
 
     # Create the pool + create user for the pool
@@ -156,14 +160,15 @@ def create_cluster(
     if wait:
         util.wait_for_master_to_be_ready(pool_id)
 
-        if username is not None and password is not None:
-            create_user(pool_id, username, password)
+        if username is not None:
+            create_user(pool_id, username, password, ssh_key)
 
 
 def create_user(
         cluster_id: str,
         username: str,
-        password: str):
+        password: str = None,
+        ssh_key: str = None) -> str:
     """
         Create a cluster user
         :param cluster_id: id of the spark cluster
@@ -171,6 +176,12 @@ def create_user(
         :param password: password of the user to add
     """
     batch_client = azure_api.get_batch_client()
+    if password is None:
+        ssh_key = ssh.get_user_public_key(ssh_key)
+
+    if not password and not ssh_key:
+        raise InvalidUserCredentialsError(
+            "Cannot add user to cluster. Need to provide a ssh public key or password.")
 
     # Create new ssh user for the master node
     batch_client.compute_node.add_user(
@@ -180,8 +191,14 @@ def create_user(
             username,
             is_admin=True,
             password=password,
+            ssh_public_key=ssh_key,
             expiry_time=datetime.now() + timedelta(days=365)))
 
+    return (
+        '*' * len(password) if password else None,
+        ssh_key,
+    )
+
 
 def get_cluster_details(cluster_id: str):
     """
@@ -294,11 +311,7 @@ def delete_cluster(cluster_id: str):
         print("Deleting cluster {0}".format(cluster_id))
 
 
-class ClusterNotReadyError(Exception):
-    pass
-
-
-def ssh(
+def ssh_in_master(
         cluster_id: str,
         username: str=None,
         masterui: str=None,
@@ -333,9 +346,12 @@ def ssh(
 
     ssh_command = CommandBuilder('ssh')
 
-    ssh_command.add_option("-L", "{0}:localhost:{1}".format(masterui, constants.SPARK_MASTER_UI_PORT), enable=bool(masterui))
-    ssh_command.add_option("-L", "{0}:localhost:{1}".format(webui, constants.SPARK_WEBUI_PORT), enable=bool(webui))
-    ssh_command.add_option("-L", "{0}:localhost:{1}".format(jupyter, constants.SPARK_JUPYTER_PORT), enable=bool(jupyter))
+    ssh_command.add_option("-L", "{0}:localhost:{1}".format(
+        masterui, constants.SPARK_MASTER_UI_PORT), enable=bool(masterui))
+    ssh_command.add_option(
+        "-L", "{0}:localhost:{1}".format(webui, constants.SPARK_WEBUI_PORT), enable=bool(webui))
+    ssh_command.add_option("-L", "{0}:localhost:{1}".format(
+        jupyter, constants.SPARK_JUPYTER_PORT), enable=bool(jupyter))
 
     if ports is not None:
         for port in ports:
@@ -343,7 +359,8 @@ def ssh(
                 "-L", "{0}:localhost:{1}".format(port[0], port[1]))
 
     user = username if username is not None else '<username>'
-    ssh_command.add_argument("{0}@{1} -p {2}".format(user, master_node_ip, master_node_port))
+    ssh_command.add_argument(
+        "{0}@{1} -p {2}".format(user, master_node_ip, master_node_port))
 
     command = ssh_command.to_str()
     ssh_command_array = command.split()

dtde/config.py

@@ -1,5 +1,5 @@
-from . import constants
 import os
+from . import constants
 
 try:
     import configparser
@@ -8,16 +8,19 @@ except ImportError:
 
 global_config = None
 
-def load_config():
+
+def load_config(path: str=constants.DEFAULT_CONFIG_PATH):
     """
         Loads the config file at the root of the repository(configuration.cfg)
     """
     global global_config
-    if not os.path.isfile(constants.CONFIG_PATH):
-        raise Exception("Configuration file doesn't exists at {0}".format(constants.CONFIG_PATH))
+    if not os.path.isfile(path):
+        raise Exception(
+            "Configuration file doesn't exists at {0}".format(path))
 
     global_config = configparser.ConfigParser()
-    global_config.read(constants.CONFIG_PATH)
+    global_config.read(path)
+
 
 def get() -> configparser.ConfigParser:
     if not global_config:

dtde/constants.py

@@ -17,7 +17,7 @@ ROOT_PATH = os.path.normpath(os.path.join(os.path.dirname(__file__), '..'))
 """
     Path to the configuration file
 """
-CONFIG_PATH = os.path.join(ROOT_PATH, 'configuration.cfg')
+DEFAULT_CONFIG_PATH = os.path.join(os.getcwd(), 'configuration.cfg')
 
 """
     Key of the metadata entry for the pool that is used to store the master node id

dtde/core/__init__.py

@@ -1,3 +1,4 @@
+from . import ssh
 from .command_builder import CommandBuilder
 
-__all__ = ["CommandBuilder"]
+__all__ = ["CommandBuilder", "ssh"]

dtde/core/ssh.py

@@ -0,0 +1,37 @@
+import os
+from dtde import config
+
+
+def get_user_public_key(key_or_path: str = None) -> str:
+    """
+        Return the ssh key.
+        It will first check if the given argument is a ssh key or a path to one
+        otherwise will check the configuration file.
+    """
+    if not key_or_path:
+        global_config = config.get()
+
+        if not global_config.has_option("Default", "ssh_pub_key"):
+            return None
+
+        key_or_path = global_config.get("Default", "ssh_pub_key")
+
+    if not key_or_path:
+        return None
+
+    key = None
+    if os.path.isfile(os.path.expanduser(key_or_path)):
+        key = __read_ssh_key_from_file(key_or_path)
+    else:
+        key = key_or_path
+
+    return key
+
+
+def __read_ssh_key_from_file(path: str) -> str:
+    """
+        Read the content of the given file
+    """
+    with open(os.path.expanduser(path), 'r') as content_file:
+        content = content_file.read()
+        return content

dtde/error.py

@@ -0,0 +1,6 @@
+class InvalidUserCredentialsError(Exception):
+    pass
+
+
+class ClusterNotReadyError(Exception):
+    pass

dtde/spark/cli/cluster_add_user.py

@@ -8,18 +8,29 @@ def setup_parser(parser: argparse.ArgumentParser):
                         help='The unique id of your spark cluster')
     parser.add_argument('-u', '--username',
                         help='The usernameto access your spark cluster\'s head node')
-    parser.add_argument('-p', '--password',
-                        help='The password to access your spark cluster\'s head node')
-    parser.set_defaults(username="admin", password="pass123!")
+
+    auth_group = parser.add_mutually_exclusive_group()
+    auth_group.add_argument('-p', '--password',
+                            help="The password to access your spark cluster's master node. If not provided will use ssh public key.")
+    auth_group.add_argument('--ssh-key',
+                            help="The ssh public key to access your spark cluster's master node. You can also set the ssh-key in the configuration file.")
+    parser.set_defaults(username="admin")
 
 
 def execute(args: typing.NamedTuple):
     print('-------------------------------------------')
     print('spark cluster id:    {}'.format(args.cluster_id))
     print('username:            {}'.format(args.username))
-    print('password:            {}'.format(args.password))
-    print('-------------------------------------------')
-    clusterlib.create_user(
+
+    password, ssh_key = clusterlib.create_user(
         args.cluster_id,
         args.username,
-        args.password)
+        args.password,
+        args.ssh_key)
+
+    if password:
+        print('password:            {}'.format(password))
+    elif ssh_key:
+        print('ssh public key:      {}'.format(ssh_key))
+
+    print('-------------------------------------------')

dtde/spark/cli/cluster_create.py

@@ -1,7 +1,6 @@
 import argparse
 import typing
-from dtde import clusterlib, util
-
+from dtde import clusterlib
 
 def setup_parser(parser: argparse.ArgumentParser):
     parser.add_argument('--id', dest='cluster_id', required=True,
@@ -21,7 +20,9 @@ def setup_parser(parser: argparse.ArgumentParser):
     parser.add_argument('--username',
                         help='Username to access your cluster (required: --wait flag)')
     parser.add_argument('--password',
-                        help='Password to access your cluster (required: --wait flag)')
+                        help="The password to access your spark cluster's head node. If not provided will use ssh public key.")
+    parser.add_argument('--ssh-key',
+                        help="The ssh public key to access your spark cluster\'s head node. You can also set the ssh-key in the configuration file.")
     parser.add_argument('--no-wait', dest='wait', action='store_false')
     parser.add_argument('--wait', dest='wait', action='store_true')
     parser.set_defaults(wait=False, size=0, size_low_pri=0)
@@ -49,4 +50,5 @@ def execute(args: typing.NamedTuple):
         args.vm_size,
         args.username,
         args.password,
+        args.ssh_key,
         args.wait)

dtde/spark/cli/cluster_ssh.py

@@ -1,6 +1,6 @@
 import argparse
 import typing
-from dtde import clusterlib, util
+from dtde import clusterlib
 
 
 def setup_parser(parser: argparse.ArgumentParser):
@@ -33,7 +33,7 @@ def execute(args: typing.NamedTuple):
     print('-------------------------------------------')
 
     # get ssh command
-    clusterlib.ssh(
+    clusterlib.ssh_in_master(
         cluster_id=args.cluster_id,
         masterui=args.masterui,
         webui=args.webui,

dtde/util.py

@@ -369,3 +369,4 @@ def get_cluster_total_current_nodes(pool):
     Get the total number of current nodes (dedicated + low pri) in the pool
     """
     return pool.current_dedicated_nodes + pool.current_low_priority_nodes
+