dcos node ssh

cli/dcoscli/node/main.py

@@ -4,18 +4,28 @@ Usage:
     dcos node --info
     dcos node [--json]
     dcos node log [--follow --lines=N --master --slave=<slave-id>]
+    dcos node ssh [--option SSHOPT=VAL ...]
+                  [--config-file=<path>]
+                  [--user=<user>]
+                  (--master | --slave=<slave-id>)
 
 Options:
-    -h, --help            Show this screen
-    --info                Show a short description of this subcommand
-    --json                Print json-formatted nodes
-    --follow              Output data as the file grows
-    --lines=N             Output the last N lines [default: 10]
-    --master              Output the leading master's Mesos log
-    --slave=<slave-id>    Output this slave's Mesos log
-    --version             Show version
+    -h, --help              Show this screen
+    --info                  Show a short description of this subcommand
+    --json                  Print json-formatted nodes
+    --follow                Output data as the file grows
+    --lines=N               Output the last N lines [default: 10]
+    --master                Access the leading master
+    --slave=<slave-id>      Access the slave with the provided ID
+    --option SSHOPT=VAL     SSH option (see `man ssh_config`)
+    --config-file=<path>    Path to ssh config file
+    --user=<user>           SSH user [default: core]
+    --version               Show version
 """
 
+import os
+import subprocess
+
 import dcoscli
 import docopt
 from dcos import cmds, emitting, errors, mesos, util
@@ -61,6 +71,12 @@ def _cmds():
             arg_keys=['--follow', '--lines', '--master', '--slave'],
             function=_log),
 
+        cmds.Command(
+            hierarchy=['node', 'ssh'],
+            arg_keys=['--master', '--slave', '--option', '--config-file',
+                      '--user'],
+            function=_ssh),
+
         cmds.Command(
             hierarchy=['node'],
             arg_keys=['--json'],
@@ -80,7 +96,7 @@ def _info():
 
 
 def _list(json_):
-    """List dcos nodes
+    """List DCOS nodes
 
     :param json_: If true, output json.
         Otherwise, output a human readable table.
@@ -89,7 +105,7 @@ def _list(json_):
     :rtype: int
     """
 
-    client = mesos.MesosClient()
+    client = mesos.DCOSClient()
     slaves = client.get_state_summary()['slaves']
     if json_:
         emitter.publish(slaves)
@@ -148,3 +164,61 @@ def _mesos_files(master, slave_id):
         slave = mesos.get_master().slave(slave_id)
         files.append(mesos.MesosFile('/slave/log', slave=slave))
     return files
+
+
+def _ssh(master, slave, option, config_file, user):
+    """SSH into a DCOS node.  Since only the masters are definitely
+    publicly available, we first ssh into an arbitrary master, then
+    hop to the desired node.
+
+    :param master: True if the user has opted to SSH into the leading
+                   master
+    :type master: bool | None
+    :param slave: The slave ID if the user has opted to SSH into a slave
+    :type slave: str | None
+    :param option: SSH option
+    :type option: [str]
+    :param config_file: SSH config file
+    :type config_file: str | None
+    :param user: SSH user
+    :type user: str | None
+    :rtype: int
+    :returns: process return code
+
+    """
+    if not os.environ.get('SSH_AUTH_SOCK'):
+        raise DCOSException(
+            "There is no SSH_AUTH_SOCK env variable, which likely means you " +
+            "aren't running `ssh-agent`.  `dcos node ssh` depends on " +
+            "`ssh-agent` so we can safely use your private key to hop " +
+            "between nodes in your cluster.  Please run `ssh-agent`, " +
+            "then add your private key with `ssh-add`.")
+
+    master_public_ip = mesos.DCOSClient().metadata()['PUBLIC_IPV4']
+    ssh_options = ' '.join('-o {}'.format(opt) for opt in option)
+
+    if config_file:
+        ssh_config = '-F {}'.format(config_file)
+    else:
+        ssh_config = ''
+
+    if master:
+        host = 'leader.mesos'
+    else:
+        summary = mesos.DCOSClient().get_state_summary()
+        slave_obj = next((slave_ for slave_ in summary['slaves']
+                          if slave_['id'] == slave),
+                         None)
+        if slave_obj:
+            host = mesos.parse_pid(slave_obj['pid'])[1]
+        else:
+            raise DCOSException('No slave found with ID [{}]'.format(slave))
+
+    cmd = "ssh -A -t {0} {1} {2}@{3} ssh -A -t {2}@{4}".format(
+        ssh_options,
+        ssh_config,
+        user,
+        master_public_ip,
+        host)
+
+    return subprocess.call(cmd, shell=True)