Merge lp:~jimbaker/pyjuju/ssh-passthrough into lp:pyjuju

[Jim Baker]

Implements the ability to passthrough args to an underlying command, in particular to support juju ssh. This support enables using juju ssh in the following ways:

$ juju ssh mysql/0 ls /

$ juju ssh -L8080:localhost:80 wordpress/0

and so forth, any options are simply passed through, along with the ssh command.

This requires working around argparse a bit, but only through its public API. In particular, main now uses parse_known_args to see if the command module in question supports the passthrough function; it if it does, that function is used to further decorate options from any extra args. Then command execution proceeds as before.

[William Reade (fwereade)]

[0] trivial

+ "[passthru ssh options] unit_or_machine [command]"),

s/passthru/passthrough

[1] trivial

+ # workding with ssh, for example), so consume the rest and

s/workding/working

[2] less trivial

I didn't see any way to handle potential conflicts in args (specifically -o ControlMaster/ControlPath, which I think are the only ones we need to worry about at the moment). Not really a big deal, but something it might be worth logging a "low" bug against, to make it clear that the limitation is known but not a priority?

anyway, +1

review: Approve

[Jim Baker (jimbaker)]

> [0] trivial
>
> + "[passthru ssh options] unit_or_machine [command]"),
>
> s/passthru/passthrough

Thanks, let's keep it consistent. (Both seem to be used, but passthrough more so. Or perhaps pass-through.)

>
> [1] trivial
>
> + # workding with ssh, for example), so consume the rest and
>
> s/workding/working

Ack

>
> [2] less trivial
>
> I didn't see any way to handle potential conflicts in args (specifically -o
> ControlMaster/ControlPath, which I think are the only ones we need to worry
> about at the moment). Not really a big deal, but something it might be worth
> logging a "low" bug against, to make it clear that the limitation is known but
> not a priority?

Indeed, it's possible for the user to escape the walled garden and use ssh flags in bad ways as well as good. Passthrough flags are presented later to ssh, and can thereby override juju's default usage. In the particular example you cite, this enables the user to know what is best for a given scenario. For example, multiplexing a connection would not be so ideal if there's significant data transfer. Until now this was a moot point, but with tunnel support, this need could actually arise, and it can be simply solved by passing in -o "ControlMaster no".

I don't think it's a low bug, just simply a matter of adding to the documentation of this command.

>
> anyway, +1

Thanks!

[Jim Baker (jimbaker)]

> I don't think it's a low bug, just simply a matter of adding to the
> documentation of this command.

To be precise: it should be fixed in the context of this branch.

[Jim Baker (jimbaker)]

>
> > I don't think it's a low bug, just simply a matter of adding to the
> > documentation of this command.
>
> To be precise: it should be fixed in the context of this branch.

irc discussion:

<jimbaker> fwereade, i think i'm going to change my mind on the doc stuff related to juju ssh. the command itself is not documented. really we need to incorporate more detailed docs on each subcommand in juju
<jimbaker> fwereade, so that should be done in a separate branch/bug
<fwereade> jimbaker, sounds sensible to me

[lp:~jimbaker/pyjuju/ssh-passthrough]

410. By Jim Baker on 2011-10-20

Fix doc typo

411. By Jim Baker on 2011-10-20

Merged trunk

412. By Jim Baker on 2011-10-20

Addressed review points

[Kapil Thangavelu (hazmat)]

Nice work jim. lgtm, a minor

[0]

+def parse_passthrough_args(args, flags_taking_arg=frozenset()):

the frozenset seems like overkill when an empty tuple suffices.

[1]

re the same function.

Is supporting long options just a matter of an additional conditional
clause for peeked.startswith("--") ? If so it seems better to just
implement support for it now.

[2] re the conflict william brought up, experimenting with the ssh cli it doesn't actually appear to be a problem in practice the user specified option will take precedence as its specified last.

review: Approve

[Jim Baker (jimbaker)]

> Nice work jim. lgtm, a minor
>
> [0]
>
> +def parse_passthrough_args(args, flags_taking_arg=frozenset()):
>
> the frozenset seems like overkill when an empty tuple suffices.

Sure, definite overkill.

>
> [1]
>
> re the same function.
>
> Is supporting long options just a matter of an additional conditional
> clause for peeked.startswith("--") ? If so it seems better to just
> implement support for it now.

It's not quite so simple, and so far the universe of commands that need this passthrough support is as far as I know is just ssh/scp. So I'd suggest punting that support until we identify them.
>
>
> [2] re the conflict william brought up, experimenting with the ssh cli it
> doesn't actually appear to be a problem in practice the user specified option
> will take precedence as its specified last.

My experience too.

=== modified file 'juju/control/__init__.py'
--- juju/control/__init__.py	2011-10-10 20:20:20 +0000
+++ juju/control/__init__.py	2011-10-20 17:26:23 +0000
@@ -4,6 +4,7 @@
 import zookeeper
 
 from .command import Commander
+from .utils import ParseError
 from juju.environment.config import EnvironmentsConfig
 
 import add_relation
@@ -106,13 +107,15 @@
 
     for module in subcommands:
         configure_subparser = getattr(module, "configure_subparser", None)
+        passthrough = getattr(module, "passthrough", None)
         if configure_subparser:
             sub_parser = configure_subparser(subparsers)
         else:
             sub_parser = subparsers.add_parser(
                 module.__name__.split('.')[-1], help=module.command.__doc__)
-        sub_parser.set_defaults(command=Commander(module.command),
-                                parser=sub_parser)
+        sub_parser.set_defaults(
+            command=Commander(module.command, passthrough=passthrough),
+            parser=sub_parser)
 
     return parser
 
@@ -157,7 +160,20 @@
         prog="juju",
         description="juju cloud orchestration admin")
     parser.set_defaults(environments=env_config, log=log)
-    options = parser.parse_args(args)
+
+    # Some commands, like juju ssh, do a further parse on options by
+    # delegating to another command (such as the underlying ssh). But
+    # first need to parse nonstrictly all args to even determine what
+    # command is even being used.
+    options, extra = parser.parse_known_args(args)
+    if options.command.passthrough:
+        try:
+            # Augments options with subparser specific passthrough parsing
+            options.command.passthrough(options, extra)
+        except ParseError, e:
+            options.parser.error(str(e))
+    else:
+        # Otherwise, do be strict
+        options = parser.parse_args(args)
     setup_logging(options)
-
     options.command(options)
 
=== modified file 'juju/control/command.py'
--- juju/control/command.py	2011-09-15 18:50:23 +0000
+++ juju/control/command.py	2011-10-20 17:26:23 +0000
@@ -19,12 +19,13 @@
 
     """
 
-    def __init__(self, callback):
+    def __init__(self, callback, passthrough=False):
         if not callable(callback):
             raise ValueError(
                 "Commander callback argument must be a callable")
 
         self.callback = callback
+        self.passthrough = passthrough
         self.options = None
         self.exit_code = 0
 
 
=== modified file 'juju/control/ssh.py'
--- juju/control/ssh.py	2011-09-25 14:05:19 +0000
+++ juju/control/ssh.py	2011-10-20 17:26:23 +0000
@@ -1,8 +1,10 @@
+from argparse import RawDescriptionHelpFormatter
 import os
 
 from twisted.internet.defer import inlineCallbacks, returnValue
 
-from juju.control.utils import get_environment
+from juju.control.utils import (
+    get_environment, parse_passthrough_args, ParseError)
 from juju.state.errors import (
     ServiceUnitStateMachineNotAssigned, MachineStateNotFound)
 from juju.state.machine import MachineStateManager
@@ -11,14 +13,34 @@
 
 
 def configure_subparser(subparsers):
-    sub_parser = subparsers.add_parser("ssh", help=command.__doc__)
-    sub_parser.add_argument(
-        "unit_or_machine", help="Name of unit or machine")
+    sub_parser = subparsers.add_parser(
+        "ssh",
+        help=command.__doc__,
+        usage=("%(prog)s [-h] [--environment ENVIRONMENT] "
+         "[passthrough ssh options] unit_or_machine [command]"),
+        formatter_class=RawDescriptionHelpFormatter,
+        description=(
+            "positional arguments:\n"
+            "  unit_or_machine       Name of unit or machine\n"
+            "  [command]             Optional command to run on machine"))
     sub_parser.add_argument(
         "--environment", "-e", help="Environment to operate on.")
     return sub_parser
 
 
+def passthrough(options, extra):
+    """Second parsing phase to parse `extra` to passthrough to ssh itself.
+
+    Partitions into flags, unit_or_machine, and optional ssh command.
+    """
+    flags, positional = parse_passthrough_args(extra, "bcDeFIiLlmOopRSWw")
+    if not positional:
+        raise ParseError("too few arguments")
+    options.ssh_flags = flags
+    options.unit_or_machine = positional.pop(0)
+    options.ssh_command = positional  # if any
+
+
 @inlineCallbacks
 def get_ip_address_for_unit(client, provider, unit_name):
     manager = ServiceStateManager(client)
@@ -41,13 +63,15 @@
     returnValue((provider_machine.dns_name, machine_state))
 
 
-def open_ssh(ip_address):
+def open_ssh(flags, ip_address, ssh_command):
     # XXX - TODO - Might be nice if we had the ability to get the user's
     # private key path and utilize it here, ie the symmetric end to
     # get user public key.
     args = ["ssh"]
     args.extend(prepare_ssh_sharing())
+    args.extend(flags)
     args.extend(["ubuntu@%s" % ip_address])
+    args.extend(ssh_command)
     os.execvp("ssh", args)
 
 
@@ -58,10 +82,9 @@
     environment = get_environment(options)
     provider = environment.get_machine_provider()
     client = yield provider.connect()
-
     label = machine = unit = None
 
-    # First check if its a juju machine id
+    # First check if it's a juju machine id
     if options.unit_or_machine.isdigit():
         options.log.debug(
             "Fetching machine address using juju machine id.")
@@ -69,7 +92,7 @@
             client, provider, options.unit_or_machine)
         machine.get_ip_address = get_ip_address_for_machine
         label = "machine"
-    # Next check if its a unit
+    # Next check if it's a unit
     elif "/" in options.unit_or_machine:
         options.log.debug(
             "Fetching machine address using unit name.")
@@ -100,4 +123,4 @@
 
     options.log.info("Connecting to %s %s at %s",
                      label, options.unit_or_machine, ip_address)
-    open_ssh(ip_address)
+    open_ssh(options.ssh_flags, ip_address, options.ssh_command)
 
=== modified file 'juju/control/tests/test_ssh.py'
--- juju/control/tests/test_ssh.py	2011-09-25 14:05:19 +0000
+++ juju/control/tests/test_ssh.py	2011-10-20 17:26:23 +0000
@@ -58,7 +58,7 @@
     @inlineCallbacks
     def test_shell_with_unit(self):
         """
-        'juju shell mysql/0' will execute ssh against the machine
+        'juju ssh mysql/0' will execute ssh against the machine
         hosting the unit.
         """
         mock_environment = self.mocker.patch(Environment)
@@ -92,6 +92,46 @@
             self.output.getvalue())
 
     @inlineCallbacks
+    def test_passthrough_args(self):
+        """Verify that args are passed through to the underlying ssh command.
+
+        For example, something like the following command should be valid::
+
+          $ juju ssh -L8080:localhost:8080 -o "ConnectTimeout 60" mysql/0 ls /
+        """
+        mock_environment = self.mocker.patch(Environment)
+        mock_environment.get_machine_provider()
+        self.mocker.result(self.provider)
+
+        mock_exec = self.mocker.replace(os.execvp)
+        mock_exec("ssh", [
+            "ssh",
+            "-o",
+            "ControlPath " + self.tmp_home + "/.juju/ssh/master-%r@%h:%p",
+            "-o", "ControlMaster no",
+            "-L8080:localhost:8080", "-o", "ConnectTimeout 60",
+            "[e-mail]", "ls *"])
+
+        # Track unwanted calls:
+        calls = []
+        mock_exec(ARGS, KWARGS)
+        self.mocker.count(0, None)
+        self.mocker.call(lambda *args, **kwargs: calls.append((args, kwargs)))
+
+        finished = self.setup_cli_reactor()
+        self.mocker.replay()
+        yield self.unit.connect_agent()
+
+        main(["ssh", "-L8080:localhost:8080", "-o", "ConnectTimeout 60",
+              self.unit.unit_name, "ls *"])
+        yield finished
+
+        self.assertEquals(calls, [])
+        self.assertIn(
+            "Connecting to unit mysql/0 at mysql-0.example.com",
+            self.output.getvalue())
+
+    @inlineCallbacks
     def test_shell_with_unit_and_unconnected_unit_agent(self):
         """If a unit doesn't have a connected unit agent,
         the ssh command will wait till one exists before connecting.
@@ -269,7 +309,7 @@
     @inlineCallbacks
     def test_shell_with_machine_id(self):
         """
-        'juju shell <machine_id>' will execute ssh against the machine
+        'juju ssh <machine_id>' will execute ssh against the machine
         with the corresponding id.
         """
         mock_environment = self.mocker.patch(Environment)
@@ -338,3 +378,25 @@
         yield finished
 
         self.assertIn("Machine 1 was not found", self.stderr.getvalue())
+
+
+class ParseErrorsTest(ServiceStateManagerTestBase, ControlToolTest):
+
+    @inlineCallbacks
+    def setUp(self):
+        yield super(ParseErrorsTest, self).setUp()
+        config = {
+            "environments": {"firstenv": {"type": "dummy"}}}
+
+        self.write_config(dump(config))
+        self.config.load()
+        self.stderr = self.capture_stream("stderr")
+
+    def test_passthrough_args_parse_error(self):
+        """Verify that bad passthrough args will get an argparse error."""
+        e = self.assertRaises(
+            SystemExit, main, ["ssh", "-L", "mysql/0"])
+        self.assertEqual(e.code, 2)
+        self.assertIn("juju ssh: error: too few arguments",
+                      self.stderr.getvalue())
+
 
=== modified file 'juju/control/tests/test_utils.py'
--- juju/control/tests/test_utils.py	2011-09-15 19:24:47 +0000
+++ juju/control/tests/test_utils.py	2011-10-20 17:26:23 +0000
@@ -2,7 +2,8 @@
 
 from juju.environment.errors import EnvironmentsConfigError
 from juju.environment.config import EnvironmentsConfig
-from juju.control.utils import get_environment
+from juju.control.utils import (
+    get_environment, parse_passthrough_args, ParseError)
 from juju.control.tests.common import ControlToolTest
 
 
@@ -65,3 +66,49 @@
             options)
 
         self.assertIn("Invalid environment 'volcano'", str(error))
+
+
+class ParsePassthroughArgsTest(ControlToolTest):
+
+    def test_parse_typical_ssh(self):
+        """Verify that flags and positional args are properly partitioned."""
+        ssh_flags = "bcDeFIiLlmOopRSWw"
+        self.assertEqual(
+            parse_passthrough_args(
+                ["-L8080:localhost:80", "-o", "Volume 11", "mysql/0", "ls a*"],
+                ssh_flags),
+            (["-L8080:localhost:80", "-o", "Volume 11"], ["mysql/0", "ls a*"]))
+
+        self.assertEqual(
+            parse_passthrough_args(
+                ["-L8080:localhost:80", "-aC26", "0", "foobar", "do", "123"],
+                ssh_flags),
+            (["-L8080:localhost:80", "-aC26"], ["0", "foobar", "do", "123"]))
+
+        self.assertEqual(
+            parse_passthrough_args(["mysql/0"], ssh_flags),
+            ([], ["mysql/0"]))
+
+        self.assertEqual(
+            parse_passthrough_args(
+                ["mysql/0", "command", "-L8080:localhost:80"], ssh_flags),
+            ([], ["mysql/0", "command", "-L8080:localhost:80"]))
+
+    def test_parse_flag_taking_args(self):
+        """Verify that arg-taking flags properly combine with args"""
+        # some sample flags, from the ssh command
+        ssh_flags = "bcDeFIiLlmOopRSWw"
+
+        for flag in ssh_flags:
+            # This flag properly combines, either of the form -Xabc or -X abc
+            self.assertEqual(
+                parse_passthrough_args(
+                    ["-" + flag + "XYZ", "-1X", "-" + flag, "XYZ", "mysql/0"],
+                    ssh_flags),
+                (["-" + flag + "XYZ", "-1X", "-" + flag, "XYZ"], ["mysql/0"]))
+
+            # And requires that it is combined
+            e = self.assertRaises(
+                ParseError,
+                parse_passthrough_args, ["-" + flag], ssh_flags)
+            self.assertEqual(str(e), "argument -%s: expected one argument" % flag)
 
=== modified file 'juju/control/utils.py'
--- juju/control/utils.py	2011-09-15 18:50:23 +0000
+++ juju/control/utils.py	2011-10-20 17:26:23 +0000
@@ -1,3 +1,5 @@
+from itertools import tee
+
 from juju.environment.errors import EnvironmentsConfigError
 
 
@@ -9,3 +11,54 @@
     elif environment is None:
         environment = options.environments.get_default()
     return environment
+
+
+class ParseError(Exception):
+    """Used to support returning custom parse errors in passthrough parsing.
+
+    Enables similar support to what is seen in argparse, without using its
+    internals.
+    """
+
+
+def parse_passthrough_args(args, flags_taking_arg=frozenset()):
+    """Scans left to right, partitioning flags and positional args.
+
+    :param args: Unparsed args from argparse
+    :param flags_taking_arg: One character flags that combine
+       with arguments.
+    :return: tuple of flags and positional args
+    :raises: :class:`juju.control.utils.ParseError`
+
+    TODO May need to support long options for other passthrough commands.
+    """
+    args = iter(args)
+    flags_taking_arg = set(flags_taking_arg)
+    flags = []
+    positional = []
+    while True:
+        args, peek_args = tee(args)
+        try:
+            peeked = peek_args.next()
+        except StopIteration:
+            break
+        if peeked.startswith("-"):
+            flags.append(args.next())
+            # Only need to consume the next arg if the flag both takes
+            # an arg (say -L) and then it has an extra arg following
+            # (8080:localhost:80), rather than being combined, such as
+            # -L8080:localhost:80
+            if len(peeked) == 2 and peeked[1] in flags_taking_arg:
+                try:
+                    flags.append(args.next())
+                except StopIteration:
+                    raise ParseError(
+                        "argument -%s: expected one argument" % peeked[1])
+        else:
+            # At this point no more flags for the command itself (more
+            # can follow after the first positional arg, as seen in
+            # working with ssh, for example), so consume the rest and
+            # stop parsing options
+            positional = list(args)
+            break
+    return flags, positional