twisted/conch/scripts/conch.py

   1 # -*- test-case-name: twisted.conch.test.test_conch -*-
   2 #
   3 # Copyright (c) Twisted Matrix Laboratories.
   4 # See LICENSE for details.
   5
   6 #
   7 # $Id: conch.py,v 1.65 2004/03/11 00:29:14 z3p Exp $
   8
   9 #""" Implementation module for the `conch` command.
  10 #"""
  11 from twisted.conch.client import connect, default, options
  12 from twisted.conch.error import ConchError
  13 from twisted.conch.ssh import connection, common
  14 from twisted.conch.ssh import session, forwarding, channel
  15 from twisted.internet import reactor, stdio, task
  16 from twisted.python import log, usage
  17
  18 import os, sys, getpass, struct, tty, fcntl, signal
  19
  20 class ClientOptions(options.ConchOptions):
  21
  22     synopsis = """Usage:   conch [options] host [command]
  23 """
  24     longdesc = ("conch is a SSHv2 client that allows logging into a remote "
  25                 "machine and executing commands.")
  26
  27     optParameters = [['escape', 'e', '~'],
  28                       ['localforward', 'L', None, 'listen-port:host:port   Forward local port to remote address'],
  29                       ['remoteforward', 'R', None, 'listen-port:host:port   Forward remote port to local address'],
  30                      ]
  31
  32     optFlags = [['null', 'n', 'Redirect input from /dev/null.'],
  33                  ['fork', 'f', 'Fork to background after authentication.'],
  34                  ['tty', 't', 'Tty; allocate a tty even if command is given.'],
  35                  ['notty', 'T', 'Do not allocate a tty.'],
  36                  ['noshell', 'N', 'Do not execute a shell or command.'],
  37                  ['subsystem', 's', 'Invoke command (mandatory) as SSH2 subsystem.'],
  38                 ]
  39
  40     compData = usage.Completions(
  41         mutuallyExclusive=[("tty", "notty")],
  42         optActions={
  43             "localforward": usage.Completer(descr="listen-port:host:port"),
  44             "remoteforward": usage.Completer(descr="listen-port:host:port")},
  45         extraActions=[usage.CompleteUserAtHost(),
  46                       usage.Completer(descr="command"),
  47                       usage.Completer(descr="argument", repeat=True)]
  48         )
  49
  50     localForwards = []
  51     remoteForwards = []
  52
  53     def opt_escape(self, esc):
  54         "Set escape character; ``none'' = disable"
  55         if esc == 'none':
  56             self['escape'] = None
  57         elif esc[0] == '^' and len(esc) == 2:
  58             self['escape'] = chr(ord(esc[1])-64)
  59         elif len(esc) == 1:
  60             self['escape'] = esc
  61         else:
  62             sys.exit("Bad escape character '%s'." % esc)
  63
  64     def opt_localforward(self, f):
  65         "Forward local port to remote address (lport:host:port)"
  66         localPort, remoteHost, remotePort = f.split(':') # doesn't do v6 yet
  67         localPort = int(localPort)
  68         remotePort = int(remotePort)
  69         self.localForwards.append((localPort, (remoteHost, remotePort)))
  70
  71     def opt_remoteforward(self, f):
  72         """Forward remote port to local address (rport:host:port)"""
  73         remotePort, connHost, connPort = f.split(':') # doesn't do v6 yet
  74         remotePort = int(remotePort)
  75         connPort = int(connPort)
  76         self.remoteForwards.append((remotePort, (connHost, connPort)))
  77
  78     def parseArgs(self, host, *command):
  79         self['host'] = host
  80         self['command'] = ' '.join(command)
  81
  82 # Rest of code in "run"
  83 options = None
  84 conn = None
  85 exitStatus = 0
  86 old = None
  87 _inRawMode = 0
  88 _savedRawMode = None
  89
  90 def run():
  91     global options, old
  92     args = sys.argv[1:]
  93     if '-l' in args: # cvs is an idiot
  94         i = args.index('-l')
  95         args = args[i:i+2]+args
  96         del args[i+2:i+4]
  97     for arg in args[:]:
  98         try:
  99             i = args.index(arg)
 100             if arg[:2] == '-o' and args[i+1][0]!='-':
 101                 args[i:i+2] = [] # suck on it scp
 102         except ValueError:
 103             pass
 104     options = ClientOptions()
 105     try:
 106         options.parseOptions(args)
 107     except usage.UsageError, u:
 108         print 'ERROR: %s' % u
 109         options.opt_help()
 110         sys.exit(1)
 111     if options['log']:
 112         if options['logfile']:
 113             if options['logfile'] == '-':
 114                 f = sys.stdout
 115             else:
 116                 f = file(options['logfile'], 'a+')
 117         else:
 118             f = sys.stderr
 119         realout = sys.stdout
 120         log.startLogging(f)
 121         sys.stdout = realout
 122     else:
 123         log.discardLogs()
 124     doConnect()
 125     fd = sys.stdin.fileno()
 126     try:
 127         old = tty.tcgetattr(fd)
 128     except:
 129         old = None
 130     try:
 131         oldUSR1 = signal.signal(signal.SIGUSR1, lambda *a: reactor.callLater(0, reConnect))
 132     except:
 133         oldUSR1 = None
 134     try:
 135         reactor.run()
 136     finally:
 137         if old:
 138             tty.tcsetattr(fd, tty.TCSANOW, old)
 139         if oldUSR1:
 140             signal.signal(signal.SIGUSR1, oldUSR1)
 141         if (options['command'] and options['tty']) or not options['notty']:
 142             signal.signal(signal.SIGWINCH, signal.SIG_DFL)
 143     if sys.stdout.isatty() and not options['command']:
 144         print 'Connection to %s closed.' % options['host']
 145     sys.exit(exitStatus)
 146
 147 def handleError():
 148     from twisted.python import failure
 149     global exitStatus
 150     exitStatus = 2
 151     reactor.callLater(0.01, _stopReactor)
 152     log.err(failure.Failure())
 153     raise
 154
 155 def _stopReactor():
 156     try:
 157         reactor.stop()
 158     except: pass
 159
 160 def doConnect():
 161 #    log.deferr = handleError # HACK
 162     if '@' in options['host']:
 163         options['user'], options['host'] = options['host'].split('@',1)
 164     if not options.identitys:
 165         options.identitys = ['~/.ssh/id_rsa', '~/.ssh/id_dsa']
 166     host = options['host']
 167     if not options['user']:
 168         options['user'] = getpass.getuser()
 169     if not options['port']:
 170         options['port'] = 22
 171     else:
 172         options['port'] = int(options['port'])
 173     host = options['host']
 174     port = options['port']
 175     vhk = default.verifyHostKey
 176     uao = default.SSHUserAuthClient(options['user'], options, SSHConnection())
 177     connect.connect(host, port, options, vhk, uao).addErrback(_ebExit)
 178
 179 def _ebExit(f):
 180     global exitStatus
 181     if hasattr(f.value, 'value'):
 182         s = f.value.value
 183     else:
 184         s = str(f)
 185     exitStatus = "conch: exiting with error %s" % f
 186     reactor.callLater(0.1, _stopReactor)
 187
 188 def onConnect():
 189 #    if keyAgent and options['agent']:
 190 #        cc = protocol.ClientCreator(reactor, SSHAgentForwardingLocal, conn)
 191 #        cc.connectUNIX(os.environ['SSH_AUTH_SOCK'])
 192     if hasattr(conn.transport, 'sendIgnore'):
 193         _KeepAlive(conn)
 194     if options.localForwards:
 195         for localPort, hostport in options.localForwards:
 196             s = reactor.listenTCP(localPort,
 197                         forwarding.SSHListenForwardingFactory(conn,
 198                             hostport,
 199                             SSHListenClientForwardingChannel))
 200             conn.localForwards.append(s)
 201     if options.remoteForwards:
 202         for remotePort, hostport in options.remoteForwards:
 203             log.msg('asking for remote forwarding for %s:%s' %
 204                     (remotePort, hostport))
 205             conn.requestRemoteForwarding(remotePort, hostport)
 206         reactor.addSystemEventTrigger('before', 'shutdown', beforeShutdown)
 207     if not options['noshell'] or options['agent']:
 208         conn.openChannel(SSHSession())
 209     if options['fork']:
 210         if os.fork():
 211             os._exit(0)
 212         os.setsid()
 213         for i in range(3):
 214             try:
 215                 os.close(i)
 216             except OSError, e:
 217                 import errno
 218                 if e.errno != errno.EBADF:
 219                     raise
 220
 221 def reConnect():
 222     beforeShutdown()
 223     conn.transport.transport.loseConnection()
 224
 225 def beforeShutdown():
 226     remoteForwards = options.remoteForwards
 227     for remotePort, hostport in remoteForwards:
 228         log.msg('cancelling %s:%s' % (remotePort, hostport))
 229         conn.cancelRemoteForwarding(remotePort)
 230
 231 def stopConnection():
 232     if not options['reconnect']:
 233         reactor.callLater(0.1, _stopReactor)
 234
 235 class _KeepAlive:
 236
 237     def __init__(self, conn):
 238         self.conn = conn
 239         self.globalTimeout = None
 240         self.lc = task.LoopingCall(self.sendGlobal)
 241         self.lc.start(300)
 242
 243     def sendGlobal(self):
 244         d = self.conn.sendGlobalRequest("conch-keep-alive@twistedmatrix.com",
 245                 "", wantReply = 1)
 246         d.addBoth(self._cbGlobal)
 247         self.globalTimeout = reactor.callLater(30, self._ebGlobal)
 248
 249     def _cbGlobal(self, res):
 250         if self.globalTimeout:
 251             self.globalTimeout.cancel()
 252             self.globalTimeout = None
 253
 254     def _ebGlobal(self):
 255         if self.globalTimeout:
 256             self.globalTimeout = None
 257             self.conn.transport.loseConnection()
 258
 259 class SSHConnection(connection.SSHConnection):
 260     def serviceStarted(self):
 261         global conn
 262         conn = self
 263         self.localForwards = []
 264         self.remoteForwards = {}
 265         if not isinstance(self, connection.SSHConnection):
 266             # make these fall through
 267             del self.__class__.requestRemoteForwarding
 268             del self.__class__.cancelRemoteForwarding
 269         onConnect()
 270
 271     def serviceStopped(self):
 272         lf = self.localForwards
 273         self.localForwards = []
 274         for s in lf:
 275             s.loseConnection()
 276         stopConnection()
 277
 278     def requestRemoteForwarding(self, remotePort, hostport):
 279         data = forwarding.packGlobal_tcpip_forward(('0.0.0.0', remotePort))
 280         d = self.sendGlobalRequest('tcpip-forward', data,
 281                                    wantReply=1)
 282         log.msg('requesting remote forwarding %s:%s' %(remotePort, hostport))
 283         d.addCallback(self._cbRemoteForwarding, remotePort, hostport)
 284         d.addErrback(self._ebRemoteForwarding, remotePort, hostport)
 285
 286     def _cbRemoteForwarding(self, result, remotePort, hostport):
 287         log.msg('accepted remote forwarding %s:%s' % (remotePort, hostport))
 288         self.remoteForwards[remotePort] = hostport
 289         log.msg(repr(self.remoteForwards))
 290
 291     def _ebRemoteForwarding(self, f, remotePort, hostport):
 292         log.msg('remote forwarding %s:%s failed' % (remotePort, hostport))
 293         log.msg(f)
 294
 295     def cancelRemoteForwarding(self, remotePort):
 296         data = forwarding.packGlobal_tcpip_forward(('0.0.0.0', remotePort))
 297         self.sendGlobalRequest('cancel-tcpip-forward', data)
 298         log.msg('cancelling remote forwarding %s' % remotePort)
 299         try:
 300             del self.remoteForwards[remotePort]
 301         except:
 302             pass
 303         log.msg(repr(self.remoteForwards))
 304
 305     def channel_forwarded_tcpip(self, windowSize, maxPacket, data):
 306         log.msg('%s %s' % ('FTCP', repr(data)))
 307         remoteHP, origHP = forwarding.unpackOpen_forwarded_tcpip(data)
 308         log.msg(self.remoteForwards)
 309         log.msg(remoteHP)
 310         if self.remoteForwards.has_key(remoteHP[1]):
 311             connectHP = self.remoteForwards[remoteHP[1]]
 312             log.msg('connect forwarding %s' % (connectHP,))
 313             return SSHConnectForwardingChannel(connectHP,
 314                                             remoteWindow = windowSize,
 315                                             remoteMaxPacket = maxPacket,
 316                                             conn = self)
 317         else:
 318             raise ConchError(connection.OPEN_CONNECT_FAILED, "don't know about that port")
 319
 320 #    def channel_auth_agent_openssh_com(self, windowSize, maxPacket, data):
 321 #        if options['agent'] and keyAgent:
 322 #            return agent.SSHAgentForwardingChannel(remoteWindow = windowSize,
 323 #                                             remoteMaxPacket = maxPacket,
 324 #                                             conn = self)
 325 #        else:
 326 #            return connection.OPEN_CONNECT_FAILED, "don't have an agent"
 327
 328     def channelClosed(self, channel):
 329         log.msg('connection closing %s' % channel)
 330         log.msg(self.channels)
 331         if len(self.channels) == 1: # just us left
 332             log.msg('stopping connection')
 333             stopConnection()
 334         else:
 335             # because of the unix thing
 336             self.__class__.__bases__[0].channelClosed(self, channel)
 337
 338 class SSHSession(channel.SSHChannel):
 339
 340     name = 'session'
 341
 342     def channelOpen(self, foo):
 343         log.msg('session %s open' % self.id)
 344         if options['agent']:
 345             d = self.conn.sendRequest(self, 'auth-agent-req@openssh.com', '', wantReply=1)
 346             d.addBoth(lambda x:log.msg(x))
 347         if options['noshell']: return
 348         if (options['command'] and options['tty']) or not options['notty']:
 349             _enterRawMode()
 350         c = session.SSHSessionClient()
 351         if options['escape'] and not options['notty']:
 352             self.escapeMode = 1
 353             c.dataReceived = self.handleInput
 354         else:
 355             c.dataReceived = self.write
 356         c.connectionLost = lambda x=None,s=self:s.sendEOF()
 357         self.stdio = stdio.StandardIO(c)
 358         fd = 0
 359         if options['subsystem']:
 360             self.conn.sendRequest(self, 'subsystem', \
 361                 common.NS(options['command']))
 362         elif options['command']:
 363             if options['tty']:
 364                 term = os.environ['TERM']
 365                 winsz = fcntl.ioctl(fd, tty.TIOCGWINSZ, '12345678')
 366                 winSize = struct.unpack('4H', winsz)
 367                 ptyReqData = session.packRequest_pty_req(term, winSize, '')
 368                 self.conn.sendRequest(self, 'pty-req', ptyReqData)
 369                 signal.signal(signal.SIGWINCH, self._windowResized)
 370             self.conn.sendRequest(self, 'exec', \
 371                 common.NS(options['command']))
 372         else:
 373             if not options['notty']:
 374                 term = os.environ['TERM']
 375                 winsz = fcntl.ioctl(fd, tty.TIOCGWINSZ, '12345678')
 376                 winSize = struct.unpack('4H', winsz)
 377                 ptyReqData = session.packRequest_pty_req(term, winSize, '')
 378                 self.conn.sendRequest(self, 'pty-req', ptyReqData)
 379                 signal.signal(signal.SIGWINCH, self._windowResized)
 380             self.conn.sendRequest(self, 'shell', '')
 381             #if hasattr(conn.transport, 'transport'):
 382             #    conn.transport.transport.setTcpNoDelay(1)
 383
 384     def handleInput(self, char):
 385         #log.msg('handling %s' % repr(char))
 386         if char in ('\n', '\r'):
 387             self.escapeMode = 1
 388             self.write(char)
 389         elif self.escapeMode == 1 and char == options['escape']:
 390             self.escapeMode = 2
 391         elif self.escapeMode == 2:
 392             self.escapeMode = 1 # so we can chain escapes together
 393             if char == '.': # disconnect
 394                 log.msg('disconnecting from escape')
 395                 stopConnection()
 396                 return
 397             elif char == '\x1a': # ^Z, suspend
 398                 def _():
 399                     _leaveRawMode()
 400                     sys.stdout.flush()
 401                     sys.stdin.flush()
 402                     os.kill(os.getpid(), signal.SIGTSTP)
 403                     _enterRawMode()
 404                 reactor.callLater(0, _)
 405                 return
 406             elif char == 'R': # rekey connection
 407                 log.msg('rekeying connection')
 408                 self.conn.transport.sendKexInit()
 409                 return
 410             elif char == '#': # display connections
 411                 self.stdio.write('\r\nThe following connections are open:\r\n')
 412                 channels = self.conn.channels.keys()
 413                 channels.sort()
 414                 for channelId in channels:
 415                     self.stdio.write('  #%i %s\r\n' % (channelId, str(self.conn.channels[channelId])))
 416                 return
 417             self.write('~' + char)
 418         else:
 419             self.escapeMode = 0
 420             self.write(char)
 421
 422     def dataReceived(self, data):
 423         self.stdio.write(data)
 424
 425     def extReceived(self, t, data):
 426         if t==connection.EXTENDED_DATA_STDERR:
 427             log.msg('got %s stderr data' % len(data))
 428             sys.stderr.write(data)
 429
 430     def eofReceived(self):
 431         log.msg('got eof')
 432         self.stdio.loseWriteConnection()
 433
 434     def closeReceived(self):
 435         log.msg('remote side closed %s' % self)
 436         self.conn.sendClose(self)
 437
 438     def closed(self):
 439         global old
 440         log.msg('closed %s' % self)
 441         log.msg(repr(self.conn.channels))
 442
 443     def request_exit_status(self, data):
 444         global exitStatus
 445         exitStatus = int(struct.unpack('>L', data)[0])
 446         log.msg('exit status: %s' % exitStatus)
 447
 448     def sendEOF(self):
 449         self.conn.sendEOF(self)
 450
 451     def stopWriting(self):
 452         self.stdio.pauseProducing()
 453
 454     def startWriting(self):
 455         self.stdio.resumeProducing()
 456
 457     def _windowResized(self, *args):
 458         winsz = fcntl.ioctl(0, tty.TIOCGWINSZ, '12345678')
 459         winSize = struct.unpack('4H', winsz)
 460         newSize = winSize[1], winSize[0], winSize[2], winSize[3]
 461         self.conn.sendRequest(self, 'window-change', struct.pack('!4L', *newSize))
 462
 463
 464 class SSHListenClientForwardingChannel(forwarding.SSHListenClientForwardingChannel): pass
 465 class SSHConnectForwardingChannel(forwarding.SSHConnectForwardingChannel): pass
 466
 467 def _leaveRawMode():
 468     global _inRawMode
 469     if not _inRawMode:
 470         return
 471     fd = sys.stdin.fileno()
 472     tty.tcsetattr(fd, tty.TCSANOW, _savedMode)
 473     _inRawMode = 0
 474
 475 def _enterRawMode():
 476     global _inRawMode, _savedMode
 477     if _inRawMode:
 478         return
 479     fd = sys.stdin.fileno()
 480     try:
 481         old = tty.tcgetattr(fd)
 482         new = old[:]
 483     except:
 484         log.msg('not a typewriter!')
 485     else:
 486         # iflage
 487         new[0] = new[0] | tty.IGNPAR
 488         new[0] = new[0] & ~(tty.ISTRIP | tty.INLCR | tty.IGNCR | tty.ICRNL |
 489                             tty.IXON | tty.IXANY | tty.IXOFF)
 490         if hasattr(tty, 'IUCLC'):
 491             new[0] = new[0] & ~tty.IUCLC
 492
 493         # lflag
 494         new[3] = new[3] & ~(tty.ISIG | tty.ICANON | tty.ECHO | tty.ECHO |
 495                             tty.ECHOE | tty.ECHOK | tty.ECHONL)
 496         if hasattr(tty, 'IEXTEN'):
 497             new[3] = new[3] & ~tty.IEXTEN
 498
 499         #oflag
 500         new[1] = new[1] & ~tty.OPOST
 501
 502         new[6][tty.VMIN] = 1
 503         new[6][tty.VTIME] = 0
 504
 505         _savedMode = old
 506         tty.tcsetattr(fd, tty.TCSANOW, new)
 507         #tty.setraw(fd)
 508         _inRawMode = 1
 509
 510 if __name__ == '__main__':
 511     run()
 512