with ssh() pt3: Modifying the AST to avoid local execution

Source

Putting part 1 and part 2 together is not much more effort. We already know how to find the body of a with statement, and we already know how to compile it, transfer it, and execute it remotely. Putting them together looks like1:

    def __enter__ (self):
        file, line= traceback.extract_stack (limit=2)[0][:2]
        code= ast.parse (open (file).read ())

        found= None
        for node in ast.walk (code):
            if type (node)==ast.With:
                if node.items[0].context_expr.func.id=='ssh':
                   if node.lineno==line:
                       found= ast.Module(body=node.body)
                       break

        if found is not None:
            data= pickle.dumps (found)
            print (ast.dump (found))

            self.client= paramiko.SSHClient ()
            self.client.load_host_keys (bash ('~/.ssh/known_hosts')[0])
            self.client.connect (*self.args, **self.kwargs)

            command= '''python3 -c "import pickle
 from ast import Module, Assign, Name, Store, Call, Load, Expr
 import sys
 c= pickle.loads (sys.stdin.buffer.read (%d))
 code= compile (c, 'remote', 'exec')
 exec (code)"''' % len (data)

           (i, o, e)= self.client.exec_command (command)
           i.write (data)
           return (i, o, e)
       else:
           raise BodyNotFoundError (file, line, code)

There are two complications that arise. One is already fixed in that code: to detect from the current entering into a context (the execution of the above method __enter__()) what file and line are we being executed. This is solved in the first line with extract_stack() from the traceback module. The only difference with the original body extraction mechanism is that we also check that we're in the right line number. Just in case, there is an exception when we don't manage to find the original code.

The second complication is... well, more complicated. We successfully execute the body in the remote and we're amused that it even works. But here's the hitch: the body is also executed locally. This is annoying.

This means that we have to not only manage to find the body of the with statement to execute it remotely, we have to make sure that it is not executed locally. In other words, we have to locally replace it with innocuous code, like pass.

Luckily, ayrton is already loading and compiling the script's code before executing it. Adding a step that somehow saves the body of all with ssh() statements but also replaces them with pass should be easy. In fact, it's disappointingly easy:

class CrazyASTTransformer (ast.NodeTransformer):

    def visit_With (self, node):
        call= node.items[0].context_expr
        if call.func.id=='ssh':
            m= Module (body=node.body)
            data= pickle.dumps (m)
            s= Bytes (s=data)
            s.lineno= node.lineno
            s.col_offset= node.col_offset
            call.args.insert (0, s)

            p= Pass ()
            p.lineno= node.lineno+1
            p.col_offset= node.col_offset+4
            node.body= [p]

This time I'm using a NodeTransformer for the task. I'm simply taking the body, wrapping it around a Module, pickling that, creating a new Bytes object with that pickle and prepending it to the arguments of ssh(). On the other hand, I'm replacing the whole body with a pass statement. So:

with ssh (...):
    <body>

Becomes:

with ssh (pickle.dumps ( Module (body=Bytes (s=<body>)) ), ...):
    pass

Easy, right? Back to the context manager, its constructor now takes the pickle of the code to execute remotely as the fist argument, and the __enter__() method now does not have to look for the code anymore.

There is one more complication that I want to address in this post, so I can more or less finish with all this. paramiko's SSHClient.exec_command() method returns a sequence of 3 objects, that represent the stdin, out and err for the remote. It would be nice if we could locally refer to them so we can interact with the remote; particularly, get it's output. This means that somehow we have to manage to capture that sequence and bind it to a local name before it's too late. There is no obvious answer for this, specially because it means that I have to create a local name, or take it from somewhere, in such a way that, either it doesn't clash with the local environment, or the user expects it in a particular name.

So I more or less chose the latter. I'm extending the construct in such a way that if we write with ssh() as foo: ..., that sequence ends in foo and you can use it after the with's body. So instead of the pass statement for local execution, I want that sequence assigned to foo. For that, we we'll need a random variable that will replace foo in the as foo part, and replace pass with foo= <random_var>. It complicates things a little, but nothing really otherworldly :)

            # take the `as foo`, make it `with ssh() as <random>: foo= <random>`
            # so we can use foo locally
            local_var= node.items[0].optional_vars.id
            remote_var= random_var ()

            # ` ... as <random>`
            node.items[0].optional_vars.id= remote_var

            # add `foo= <random>` to the body
            last_lineno= node.body[-1].lineno
            col_offset= node.body[0].col_offset

            target= Name(id=local_var, ctx=Store())
            target.lineno= last_lineno+1
            target.col_offset= col_offset

            value= Name(id=remote_var, ctx=Load())
            value.lineno= last_lineno+1
            # this is a little AR :)
            value.col_offset= col_offset+len (local_var)+ len ('= ')

            ass= Assign (targets=[target], value=value)
            ass.lineno= last_lineno+1
            ass.col_offset= col_offset

            node.body= [ ass ]

Just to be clear, the final code looks like:

with ssh (pickle.dumps ( Module (body=Bytes (s=<body>)) ), ...) as <random_var>:
    foo= <random_var>

The next step is to be able to pass the locals() to the remote so it can access the local values.

  1. One note here: I finally managed to make paramiko behave in Python3, so this code is slightly different from the one in the previous post. ↩