pyssh/__init__.py

   1 """A SSH Interface class.\r
   2 \r
   3 An interface to ssh on posix systems, and plink (part of the Putty\r
   4 suite) on Win32 systems.\r
   5 \r
   6 By Rasjid Wilcox.\r
   7 Copyright (c) 2002.\r
   8 \r
   9 Version: 0.2\r
  10 Last modified 4 September 2002.\r
  11 \r
  12 Drawing on ideas from work by Julian Schaefer-Jasinski, Guido's telnetlib and\r
  13 version 0.1 of pyssh (http://pyssh.sourceforge.net) by Chuck Esterbrook.\r
  14 \r
  15 Licenced under a Python 2.2 style license.  See License.txt.\r
  16 """\r
  17 \r
  18 DEBUG_LEVEL = 0\r
  19 \r
  20 import os, getpass\r
  21 import signal    # should cause all KeyboardInterrupts to go to the main thread\r
  22                  # try for Linux, does not seem to be try under Cygwin\r
  23 import nbpipe\r
  24 import time\r
  25 \r
  26 # Constants\r
  27 SSH_PORT=806\r
  28 SSH_PATH=''\r
  29 \r
  30 CTRL_C=chr(3)\r
  31 \r
  32 READ_LAZY=0\r
  33 READ_SOME=1\r
  34 READ_ALL=2\r
  35 READ_UNTIL=3\r
  36 \r
  37 # set the path to ssh / plink, and chose the popen2 funciton to use\r
  38 if os.name=='posix':\r
  39     import fssa    # we can look for ssh-agent on posix\r
  40                    # XXX Can we on Win32/others?\r
  41     import ptyext  # if my patch gets accepted, change this to check for a\r
  42                    # sufficiently high version of python, and assign ptyext=pty\r
  43                    # if sufficient.\r
  44     sshpopen2=ptyext.popen2\r
  45     CLOSE_STR='~.'\r
  46     tp=os.popen('/usr/bin/which ssh')\r
  47     SSH_PATH=tp.read().strip()\r
  48     try:\r
  49         tp.close()\r
  50     except IOError:\r
  51         # probably no child process\r
  52         pass\r
  53     if SSH_PATH == '':\r
  54         tp=os.popen('command -v ssh')  # works in bash, ash etc, not csh etc.\r
  55         SSH_PATH=tp.read().strip()\r
  56         tp.close()\r
  57     if SSH_PATH == '':\r
  58         check = ['/usr/bin/ssh', '/usr/local/bin/ssh', '/bin/ssh']\r
  59         for item in check:\r
  60             if os.path.isfile(item):\r
  61                 SSH_PATH=item\r
  62                 break\r
  63     PORT_STR='-p '\r
  64 else:\r
  65     sshpopen2=os.popen2\r
  66     CLOSE_STR=CTRL_C        # FIX-ME: This does not work.\r
  67                             # I think I need to implement a 'kill' component\r
  68                             # to the close function using win32api.\r
  69     SSH_PATH=''\r
  70     PORT_STR='-P '\r
  71 \r
  72 class mysshError(Exception):\r
  73     """Error class for myssh."""\r
  74     pass\r
  75 \r
  76 # Helper functions\r
  77 def _prompt(prompt):\r
  78     """Print the message as the prompt for input.\r
  79     Return the text entered."""\r
  80     noecho = (prompt.lower().find('password:') >= 0) or \\r
  81         (prompt.lower().find('passphrase:') >=0)\r
  82     print """User input required for ssh connection.\r
  83     (Type Ctrl-C to abort connection.)"""\r
  84     abort = 0\r
  85     try:\r
  86         if noecho:\r
  87             response = getpass.getpass(prompt)\r
  88         else:\r
  89             response = raw_input(prompt)\r
  90     except KeyboardInterrupt:\r
  91         response = ''\r
  92         abort = 1\r
  93     return response, abort\r
  94 \r
  95 class Ssh:\r
  96     """A SSH connection class."""\r
  97     def __init__(self, username=None, host='localhost', port=None):\r
  98         """Constructor.  This does not try to connect."""\r
  99         self.debuglevel = DEBUG_LEVEL\r
 100         self.sshpath = SSH_PATH\r
 101         self.username = username\r
 102         self.host = host\r
 103         self.port = port\r
 104         self.isopen = 0\r
 105         self.sshpid = 0  # perhaps merge this with isopen\r
 106         self.old_handler = signal.getsignal(signal.SIGCHLD)\r
 107         sig_handler = signal.signal(signal.SIGCHLD, self.sig_handler)\r
 108         \r
 109     def __del__(self):\r
 110         """Destructor -- close the connection."""\r
 111         if self.isopen:\r
 112             self.close()\r
 113     \r
 114     def sig_handler(self, signum, stack):\r
 115         """ Handle SIGCHLD signal """\r
 116         if signum == signal.SIGCHLD:\r
 117             try:\r
 118                 os.waitpid(self.sshpid, 0)\r
 119             except:\r
 120                 pass\r
 121         if self.old_handler != signal.SIG_DFL:\r
 122             self.old_handler(signum, stack)\r
 123 \r
 124     def attach_agent(self, key=None):\r
 125         if os.name != 'posix':\r
 126             # only posix support at this time\r
 127             return\r
 128         if 'SSH_AUTH_SOCK' not in os.environ.keys():\r
 129             fssa.fssa(key)\r
 130 \r
 131     def set_debuglevel(self, debuglevel):\r
 132         """Set the debug level."""\r
 133         self.debuglevel = debuglevel\r
 134         \r
 135     def set_sshpath(self, sshpath):\r
 136         """Set the ssh path."""\r
 137         self.sshpath=sshpath\r
 138     \r
 139     # Low level functions\r
 140     def open(self, cmd=None):\r
 141         """Opens a ssh connection.\r
 142         \r
 143         Raises an mysshError if myssh.sshpath is not a file.\r
 144         Raises an error if attempting to open an already open connection.\r
 145         """\r
 146         self.attach_agent()\r
 147         if not os.path.isfile(self.sshpath):\r
 148             raise mysshError, \\r
 149             "Path to ssh or plink is not defined or invalid.\nsshpath='%s'" \\r
 150              % self.sshpath\r
 151         if self.isopen:\r
 152             raise mysshError, "Connection already open."\r
 153         sshargs = ''\r
 154         if self.sshpath.lower().find('plink') != -1:\r
 155             sshargs = '-ssh '\r
 156         if self.port and self.port != '':\r
 157             sshargs += PORT_STR + self.port + ' '\r
 158         if self.username and self.username !='':\r
 159             sshargs += self.username + '@'\r
 160         sshargs += self.host\r
 161         if cmd:\r
 162             sshargs += ' ' + cmd\r
 163         if self.debuglevel:\r
 164             print ">> Running %s %s." % (self.sshpath, sshargs)\r
 165         # temporary workaround until I get pid's working under win32\r
 166         if os.name == 'posix':\r
 167             self.sshin, self.sshoutblocking, self.sshpid = \\r
 168                                 sshpopen2(self.sshpath + ' ' + sshargs)\r
 169         else:\r
 170             self.sshin, self.sshoutblocking = \\r
 171                                 sshpopen2(self.sshpath + ' ' + sshargs)\r
 172         self.sshout = nbpipe.nbpipe(self.sshoutblocking)\r
 173         self.isopen = 1\r
 174         if self.debuglevel:\r
 175             print ">> ssh pid is %s." % self.sshpid\r
 176         \r
 177     def close(self, addnewline=1):\r
 178         """Close the ssh connection by closing the input and output pipes.\r
 179         Returns the closing messages.\r
 180         \r
 181         On Posix systems, by default it adds a newline before sending the\r
 182         disconnect escape sequence.   Turn this off by setting addnewline=0.\r
 183         """\r
 184         if os.name == 'posix':\r
 185             try:\r
 186                 if addnewline:\r
 187                     self.write('\n')\r
 188                 self.write(CLOSE_STR)\r
 189             except (OSError, IOError, mysshError):\r
 190                 pass\r
 191         output = self.read_lazy()\r
 192         try:\r
 193             self.sshin.close()\r
 194             self.sshoutblocking.close()\r
 195         except:\r
 196             pass\r
 197         if os.name == 'posix':\r
 198             try:\r
 199                 os.kill(self.sshpid, signal.SIGHUP)\r
 200                 os.waitpid(self.sshpid, 0)\r
 201             except:\r
 202                 pass\r
 203         self.isopen = 0\r
 204         if self.debuglevel:\r
 205             print ">> Connection closed."\r
 206         return output\r
 207         \r
 208     def write(self, text):\r
 209         """Send text to the ssh process."""\r
 210         # May block?? Probably not in practice, as ssh has a large input buffer.\r
 211         if self.debuglevel:\r
 212             print ">> Sending %s" % text\r
 213         if self.isopen:\r
 214             while len(text):\r
 215                 numtaken = os.write(self.sshin.fileno(),text)\r
 216                 if self.debuglevel:\r
 217                     print ">> %s characters taken" % numtaken\r
 218                 text = text[numtaken:]\r
 219         else:\r
 220             raise mysshError, "Attempted to write to closed connection."\r
 221     \r
 222     # There is a question about what to do with connections closed by the other\r
 223     # end.  Should write and read check for this, and force proper close?\r
 224     def read_lazy(self):\r
 225         """Lazy read from sshout.  Waits a little, but does not block."""\r
 226         return self.sshout.read_lazy()\r
 227     \r
 228     def read_some(self):\r
 229         """Always read at least one block, unless the connection is closed.\r
 230         My block."""\r
 231         if self.isopen:\r
 232             return self.sshout.read_some()\r
 233         else:\r
 234             return self.sshout.read_lazy()\r
 235         \r
 236     def read_until(self, match, timeout=None):\r
 237         """Read until a given string is encountered or until timeout.\r
 238 \r
 239         When no match is found, return whatever is available instead,\r
 240         possibly the empty string.  Raise EOFError if the connection\r
 241         is closed and no cooked data is available.\r
 242 \r
 243         """\r
 244         if self.isopen:\r
 245             return self.sshout.read_until(match, timeout)\r
 246         else:\r
 247             return self.sshout.read_lazy()    \r
 248 \r
 249     def read_all(self):\r
 250         """Reads until end of file hit.  May block."""\r
 251         if self.isopen:\r
 252             return self.sshout.read_all()\r
 253         else:\r
 254             return self.sshout.read_lazy()    \r
 255         \r
 256     # High level funcitons\r
 257     def login(self, logintext='Last login:', prompt_callback=_prompt):\r
 258         """Logs in to the ssh host.  Checks for standard prompts, and calls\r
 259         the function passed as promptcb to process them.\r
 260         Returns the login banner, or 'None' if login process aborted.\r
 261         """\r
 262         self.open()\r
 263         banner = self.read_some()\r
 264         if self.debuglevel:\r
 265             print ">> 1st banner read is: %s" % banner\r
 266         while banner.find(logintext) == -1:\r
 267             response, abort = prompt_callback(banner)\r
 268             if abort:\r
 269                 return self.close()\r
 270             self.write(response + '\n')\r
 271             banner = self.read_some()\r
 272         return banner\r
 273     \r
 274     def logout(self):\r
 275         """Logs out the session."""\r
 276         self.close()\r
 277         \r
 278     def sendcmd(self, cmd, readtype=READ_SOME, expected=None, timeout=None):\r
 279         """Sends the command 'cmd' over the ssh connection, and returns the\r
 280         result.  By default it uses read_some, which may block.\r
 281         """\r
 282         if cmd[-1] != '\n':\r
 283             cmd += '\n'\r
 284         self.write(cmd)\r
 285         if readtype == READ_ALL:\r
 286             return self.read_all()\r
 287         elif readtype == READ_LAZY:\r
 288             return self.read_lazy()\r
 289         elif readtype == READ_UNTIL and expected is not None:\r
 290             return self.read_until(expected, timeout)\r
 291         else:\r
 292             return self.read_some()\r
 293     \r
 294 def test():\r
 295     """Test routine for myssh.\r
 296     \r
 297     Usage: python myssh.py [-d] [-sshp path-to-ssh] [username@host | host] [port]\r
 298     \r
 299     Default host is localhost, default port is 22.\r
 300     """\r
 301     import sys\r
 302     debug = 0\r
 303     if sys.argv[1:] and sys.argv[1] == '-d':\r
 304         debug = 1\r
 305         del sys.argv[1]\r
 306     testsshpath = SSH_PATH\r
 307     if sys.argv[1:] and sys.argv[1] == '-sshp':\r
 308         testsshpath = sys.argv[2]\r
 309         del sys.argv[1]\r
 310         del sys.argv[1]\r
 311     testusername = None\r
 312     testhost = 'localhost'\r
 313     testport = '22'\r
 314     if sys.argv[1:]:\r
 315         testhost = sys.argv[1]\r
 316         if testhost.find('@') != -1:\r
 317             testusername, testhost = testhost.split('@')\r
 318     if sys.argv[2:]:\r
 319         testport = sys.argv[2]\r
 320         \r
 321     testcon = Ssh(testusername, testhost, testport)\r
 322     testcon.set_debuglevel(debug)\r
 323     testcon.set_sshpath(testsshpath)\r
 324     testcon.login()\r
 325     \r
 326     cmd = None\r
 327     while (cmd != 'exit') and testcon.isopen:\r
 328         cmd = raw_input("Enter command to send: ")\r
 329         print testcon.sendcmd(cmd)\r
 330     testcon.close()\r
 331 \r
 332 if __name__ == '__main__':\r
 333     test()\r