pcucontrol/transports/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 = 1\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('bash -c "type -p 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 \r
 148         if not os.path.isfile(self.sshpath):\r
 149             raise mysshError, \\r
 150             "Path to ssh or plink is not defined or invalid.\nsshpath='%s'" \\r
 151              % self.sshpath\r
 152         if self.isopen:\r
 153             raise mysshError, "Connection already open."\r
 154         sshargs = ''\r
 155         if self.sshpath.lower().find('plink') != -1:\r
 156             sshargs = '-ssh '\r
 157         if self.port and self.port != '':\r
 158             sshargs += PORT_STR + self.port + ' '\r
 159         sshargs += " -o StrictHostKeyChecking=no -o PasswordAuthentication=yes -o PubkeyAuthentication=no "\r
 160         if self.username and self.username !='':\r
 161             sshargs += self.username + '@'\r
 162         sshargs += self.host\r
 163         if cmd:\r
 164             sshargs += ' ' + cmd\r
 165         if self.debuglevel:\r
 166             print ">> Running %s %s." % (self.sshpath, sshargs)\r
 167         # temporary workaround until I get pid's working under win32\r
 168         #print sshargs\r
 169         if os.name == 'posix':\r
 170             self.sshin, self.sshoutblocking, self.sshpid = \\r
 171                                 sshpopen2(self.sshpath + ' ' + sshargs)\r
 172         else:\r
 173             self.sshin, self.sshoutblocking = \\r
 174                                 sshpopen2(self.sshpath + ' ' + sshargs)\r
 175         self.sshout = nbpipe.nbpipe(self.sshoutblocking)\r
 176         self.isopen = 1\r
 177         if self.debuglevel:\r
 178             print ">> ssh pid is %s." % self.sshpid\r
 179         \r
 180     def close(self, addnewline=1):\r
 181         """Close the ssh connection by closing the input and output pipes.\r
 182         Returns the closing messages.\r
 183         \r
 184         On Posix systems, by default it adds a newline before sending the\r
 185         disconnect escape sequence.   Turn this off by setting addnewline=0.\r
 186         """\r
 187         if os.name == 'posix':\r
 188             try:\r
 189                 if addnewline:\r
 190                     self.write('\n')\r
 191                 self.write(CLOSE_STR)\r
 192             except (OSError, IOError, mysshError):\r
 193                 pass\r
 194         output = self.read_lazy()\r
 195         try:\r
 196             self.sshin.close()\r
 197             self.sshoutblocking.close()\r
 198         except:\r
 199             pass\r
 200         if os.name == 'posix':\r
 201             try:\r
 202                 os.kill(self.sshpid, signal.SIGHUP)\r
 203                 os.waitpid(self.sshpid, 0)\r
 204             except:\r
 205                 pass\r
 206         self.isopen = 0\r
 207         if self.debuglevel:\r
 208             print ">> Connection closed."\r
 209         return output\r
 210         \r
 211     def write(self, text):\r
 212         """Send text to the ssh process."""\r
 213         # May block?? Probably not in practice, as ssh has a large input buffer.\r
 214         if self.debuglevel:\r
 215             print ">> Sending %s" % text\r
 216         if self.isopen:\r
 217             while len(text):\r
 218                 numtaken = os.write(self.sshin.fileno(),text)\r
 219                 if self.debuglevel:\r
 220                     print ">> %s characters taken" % numtaken\r
 221                 text = text[numtaken:]\r
 222         else:\r
 223             raise mysshError, "Attempted to write to closed connection."\r
 224     \r
 225     # There is a question about what to do with connections closed by the other\r
 226     # end.  Should write and read check for this, and force proper close?\r
 227     def read_lazy(self):\r
 228         """Lazy read from sshout.  Waits a little, but does not block."""\r
 229         return self.sshout.read_lazy()\r
 230     \r
 231     def read_some(self):\r
 232         """Always read at least one block, unless the connection is closed.\r
 233         My block."""\r
 234         if self.isopen:\r
 235             return self.sshout.read_some()\r
 236         else:\r
 237             return self.sshout.read_lazy()\r
 238         \r
 239     def read_until(self, match, timeout=None):\r
 240         """Read until a given string is encountered or until timeout.\r
 241 \r
 242         When no match is found, return whatever is available instead,\r
 243         possibly the empty string.  Raise EOFError if the connection\r
 244         is closed and no cooked data is available.\r
 245 \r
 246         """\r
 247         if self.isopen:\r
 248             return self.sshout.read_until(match, timeout)\r
 249         else:\r
 250             return self.sshout.read_lazy()    \r
 251 \r
 252     def read_all(self):\r
 253         """Reads until end of file hit.  May block."""\r
 254         if self.isopen:\r
 255             return self.sshout.read_all()\r
 256         else:\r
 257             return self.sshout.read_lazy()    \r
 258         \r
 259     # High level funcitons\r
 260     def login(self, logintext='Last login:', prompt_callback=_prompt):\r
 261         """Logs in to the ssh host.  Checks for standard prompts, and calls\r
 262         the function passed as promptcb to process them.\r
 263         Returns the login banner, or 'None' if login process aborted.\r
 264         """\r
 265         self.open()\r
 266         banner = self.read_some()\r
 267         if self.debuglevel:\r
 268             print ">> 1st banner read is: %s" % banner\r
 269         while banner.find(logintext) == -1:\r
 270             response, abort = prompt_callback(banner)\r
 271             if abort:\r
 272                 return self.close()\r
 273             self.write(response + '\n')\r
 274             banner = self.read_some()\r
 275         return banner\r
 276     \r
 277     def logout(self):\r
 278         """Logs out the session."""\r
 279         self.close()\r
 280         \r
 281     def sendcmd(self, cmd, readtype=READ_SOME, expected=None, timeout=None):\r
 282         """Sends the command 'cmd' over the ssh connection, and returns the\r
 283         result.  By default it uses read_some, which may block.\r
 284         """\r
 285         if cmd[-1] != '\n':\r
 286             cmd += '\n'\r
 287         self.write(cmd)\r
 288         if readtype == READ_ALL:\r
 289             return self.read_all()\r
 290         elif readtype == READ_LAZY:\r
 291             return self.read_lazy()\r
 292         elif readtype == READ_UNTIL and expected is not None:\r
 293             return self.read_until(expected, timeout)\r
 294         else:\r
 295             return self.read_some()\r
 296     \r
 297 def test():\r
 298     """Test routine for myssh.\r
 299     \r
 300     Usage: python myssh.py [-d] [-sshp path-to-ssh] [username@host | host] [port]\r
 301     \r
 302     Default host is localhost, default port is 22.\r
 303     """\r
 304     import sys\r
 305     debug = 0\r
 306     if sys.argv[1:] and sys.argv[1] == '-d':\r
 307         debug = 1\r
 308         del sys.argv[1]\r
 309     testsshpath = SSH_PATH\r
 310     if sys.argv[1:] and sys.argv[1] == '-sshp':\r
 311         testsshpath = sys.argv[2]\r
 312         del sys.argv[1]\r
 313         del sys.argv[1]\r
 314     testusername = None\r
 315     testhost = 'localhost'\r
 316     testport = '22'\r
 317     if sys.argv[1:]:\r
 318         testhost = sys.argv[1]\r
 319         if testhost.find('@') != -1:\r
 320             testusername, testhost = testhost.split('@')\r
 321     if sys.argv[2:]:\r
 322         testport = sys.argv[2]\r
 323         \r
 324     testcon = Ssh(testusername, testhost, testport)\r
 325     testcon.set_debuglevel(debug)\r
 326     testcon.set_sshpath(testsshpath)\r
 327     testcon.login()\r
 328     \r
 329     cmd = None\r
 330     while (cmd != 'exit') and testcon.isopen:\r
 331         cmd = raw_input("Enter command to send: ")\r
 332         print testcon.sendcmd(cmd)\r
 333     testcon.close()\r
 334 \r
 335 if __name__ == '__main__':\r
 336     test()\r