Scripting SSH with Python

How to Script SSH with Python

The other day Aaron Passey, our CTO, pointed out to me that I’ve written the same bit of code in different languages at each of the last 3 places I’ve worked. I figured that since this seems to be something commonly useful, yet not obviously available, I’d write something about it.

The scenario is that we have some code which wants to do remote ssh calls. Some variant of this code exists within the Isilon cluster management code, and we use it at Clustrix within our clx command line tool, as well as the database update scripts. We’re already running on a Unix box (or variant [osx, linux, etc.]) and we have access to an ssh client, so this it totally do-able from code, but not immediately obvious how.

I’ll present this in Python, but the same applies to C, or any other language. I’m aware of the Paramiko library for Python which is supposed to have support for this. That may do everything this can do – I don’t know. This is 100 lines of code, pretty easy to follow, and maybe instructive. This is all about utilizing ssh and our friendly posix primitives.  After the walkthrough, I’ve included the entire source at the bottom.

The obvious thing that we’d like to do is fork off an ssh process and read the stdin. The easy way to do this in python is with os.popen2(). This will give us back the stdin and stdout:

(sin, sout) = os.popen2(cmd)

This, however, will not work. ssh wants a psuedo tty (a pty). If it’s not running in one, it just exits. This is where the helpful python pty class comes in:

(pid, f) = pty.fork()

Now we’ve got ssh running in the right environment.  The two args we got back are important: pid is the process id of our forked ssh process, and f is a unix fileno (not to be confused with a python file handle) which is the combined stdin and stdout of the process. It’s important to remember that f isn’t something reference counted. We’re going to need to explicitly close it.

Now that we have the basic mechanism in place, let’s build us a little ssh class:

class SSH:
    def __init__(self, ip, passwd, user, port):
        self.ip = ip
        self.passwd = passwd
        self.user = user
        self.port = port

This is structured so you can create one SSH object per target box and reuse it to do different commands. We’ll also include the ability to push and pull files. Our first method will be the command handler. It will take only one argument (other than self), which will be the command to run:

 def run_cmd(self, c):
        (pid, f) = pty.fork()
        if pid == 0:
            os.execlp("ssh", "ssh", '-p %d' % self.port,
                      self.user + '@' + self.ip, c)
        else:
            return (pid, f)

As you can see, the pty fork works just like the os fork in that if the pid is 0 it means we’re the child, and if non 0 it means we’re the parent.

Since this is a raw unix fileno, the file closed condition is a little weird. Reading it will block until something is available and then return results. But when the descriptor closes it throws an os error. I’d rather be able to handle the reads just in a loop (or maybe it’s because I’m an old C programmer) so I’m going to wrap the read:

 def _read(self, f):
        x = ''
        try:
            x = os.read(f, 1024)
        except Exception, e:
            # this always fails with io error
            pass
        return x

Once we’ve got this thing forked, and can read from it, we need to get our results back out. There’s on additional thing we need to be prepared for: ssh might want to ask us some questions. We’ve all seen this before:

harmony:~$ ssh paulmini
The authenticity of host 'paulmini (10.1.2.125)' can't be established.
RSA key fingerprint is 7e:91:5d:5d:06:fe:3f:24:94:84:
c0:75:96:8c:d1:f1.
Are you sure you want to continue connecting (yes/no)?

We just want to say “yes” and move on. ssh might also ask us for a password if we don’t have host keys enabled and we’ll need to be prepared to handle that. We’ll handle all of these actions in an ssh_results method:

  def ssh_results(self, pid, f):

First, let’s initialize our output buffer:

        output = “”

Now we’ll read our first chunk and see if ssh is asking anything of us. If they want to know if we really want to continue connecting because the target isn’t in ssh/known_hosts, we’ll say yes. If they ask us for the password we’ll provide it.

got = self._read(f)
        # check for authenticity of host request
        m = re.search(“authenticity of host”, got)
        if m:
            os.write(f, ‘yesn’)
            # Read until we get ack
            while True:
                got = self._read(f)
                m = re.search(“Permanently added”, got)
                if m:
                    break

            got = self._read(f)
         # check for passwd request
        m = re.search("assword:", got)
        if m:
            # send passwd
            os.write(f, self.passwd + 'n')
            # read two lines
            tmp = self._read(f)
            tmp += self._read(f)
            m = re.search("Permission denied", tmp)
            if m:
                raise Exception("Invalid passwd")
            # passwd was accepted
            got = tmp

Preliminaries done, we can now entire our results loop:

 while got and len(got) > 0:
            output += got
            got = self._read(f)

We’ve know we’ve now ready everything because our _read returned empty. This also means that the ssh process has ended. There’s a defunct zombie process sitting there and we’re going to have to clean that up. We could handle the SIGCHLD and waitpid it, but in this case since we know the pid and we know the process is done, it’s much simpler. Also, remember that since f isn’t a reference counted Python object we’re going to need to manually clean that up:

        os.waitpid(pid, 0)
        os.close(f)
        return output

And that’s basically it. With some niceties for pushing and pulling files, and some error handling, the completed code looks like this:

#
# Remote ssh cmds
#

import pty, re, os, sys, stat, getpass

class SSHError(Exception):
    def __init__(self, value):
        self.value = value
    def __str__(self):
        return repr(self.value)
class SSH: 
    def __init__(self, ip, passwd, user, port):
        self.ip = ip
        self.passwd = passwd
        self.user = user
        self.port = port
    def run_cmd(self, c):
        (pid, f) = pty.fork()
        if pid == 0:
            os.execlp(“ssh”, “ssh”, ‘-p %d’ % self.port,
                      self.user + ‘@’ + self.ip, c)
        else:
            return (pid, f)     def push_file(self, src, dst):
        (pid, f) = pty.fork()
        if pid == 0:
            os.execlp(“scp”, “scp”, ‘-P %d’ % self.port,
                      src, self.user + ‘@’ + self.ip + ‘:’ + dst)
        else:
            return (pid, f) 


    def push_dir(self, src, dst):
        (pid, f) = pty.fork()
        if pid == 0:
            os.execlp(“scp”, “scp”, ‘-P %d’ % self.port, “-r”, src,
                     self.user + ‘@’ + self.ip + ‘:’ + dst)
        else:
            return (pid, f)

 

    def _read(self, f):
        x = ”
        try:
            x = os.read(f, 1024)
        except Exception, e:
            # this always fails with io error
            pass
        return x

    def ssh_results(self, pid, f):
        output = “”
        got = self._read(f)         # check for authenticity of host request
        m = re.search(“authenticity of host”, got)
        if m:
            os.write(f, ‘yesn’) 
            # Read until we get ack
           while True:
                got = self._read(f)
                m = re.search(“Permanently added”, got)
                if m:
                    break

            got = self._read(f)         # check for passwd request
        m = re.search(“assword:”, got)
        if m:
            # send passwd
            os.write(f, self.passwd + ‘n’)
            # read two lines
            tmp = self._read(f)
            tmp += self._read(f)
            m = re.search(“Permission denied”, tmp)
            if m:
                raise Exception(“Invalid passwd”)
            # passwd was accepted
            got = tmp
        while got and len(got) > 0:
            output += got
            got = self._read(f)
        os.waitpid(pid, 0)
        os.close(f)
        return output
    def cmd(self, c):
        (pid, f) = self.run_cmd(c)
        return self.ssh_results(pid, f)
    def push(self, src, dst):
        s = os.stat(src)
        if stat.S_ISDIR(s[stat.ST_MODE]):
            (pid, f) = self.push_dir(src, dst)
        else:
            (pid, f) = self.push_file(src, dst)
        return self.ssh_results(pid, f)
def ssh_cmd(ip, passwd, cmd, user=getpass.getuser(), port=22):
    s = SSH(ip, passwd, user, port)
    return s.cmd(cmd)
def ssh_push(ip, passwd, src, dst, user=getpass.getuser(), port=22): 
    s = SSH(ip, passwd, user, port)
    return s.push(src, dst)