root/trunk/twisted/conch/scripts/conch.py

Revision 32921, 17.7 KB (checked in by teratorn, 7 months ago)

Merge make-zshcomp-dynamic-3078-6: New tab-completion system
Author: teratorn
Reviewer: glyph, exarkun
Fixes: #3078

Deprecates t.p.zshcomp in favor of a tab-completion system in t.p.usage - completion matches are generated dynamically at tab-press time.
Line 
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#"""
11from twisted.conch.client import connect, default, options
12from twisted.conch.error import ConchError
13from twisted.conch.ssh import connection, common
14from twisted.conch.ssh import session, forwarding, channel
15from twisted.internet import reactor, stdio, task
16from twisted.python import log, usage
17
18import os, sys, getpass, struct, tty, fcntl, signal
19
20class 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"
83options = None
84conn = None
85exitStatus = 0
86old = None
87_inRawMode = 0
88_savedRawMode = None
89
90def 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
147def 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
155def _stopReactor():
156    try:
157        reactor.stop()
158    except: pass
159
160def 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
179def _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
188def 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
221def reConnect():
222    beforeShutdown()
223    conn.transport.transport.loseConnection()
224
225def beforeShutdown():
226    remoteForwards = options.remoteForwards
227    for remotePort, hostport in remoteForwards:
228        log.msg('cancelling %s:%s' % (remotePort, hostport))
229        conn.cancelRemoteForwarding(remotePort)
230
231def stopConnection():
232    if not options['reconnect']:
233        reactor.callLater(0.1, _stopReactor)
234
235class _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
259class 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
338class 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
464class SSHListenClientForwardingChannel(forwarding.SSHListenClientForwardingChannel): pass
465class SSHConnectForwardingChannel(forwarding.SSHConnectForwardingChannel): pass
466
467def _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
475def _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
510if __name__ == '__main__':
511    run()
Note: See TracBrowser for help on using the browser.