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

Revision 32921, 22.4 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_scripts -*-
2# Copyright (c) Twisted Matrix Laboratories.
3# See LICENSE for details.
4
5#
6# $Id: tkconch.py,v 1.6 2003/02/22 08:10:15 z3p Exp $
7
8""" Implementation module for the `tkconch` command.
9"""
10
11from __future__ import nested_scopes
12
13import Tkinter, tkFileDialog, tkFont, tkMessageBox, string
14from twisted.conch.ui import tkvt100
15from twisted.conch.ssh import transport, userauth, connection, common, keys
16from twisted.conch.ssh import session, forwarding, channel
17from twisted.conch.client.default import isInKnownHosts
18from twisted.internet import reactor, defer, protocol, tksupport
19from twisted.python import usage, log
20
21import os, sys, getpass, struct, base64, signal
22
23class TkConchMenu(Tkinter.Frame):
24    def __init__(self, *args, **params):
25        ## Standard heading: initialization
26        apply(Tkinter.Frame.__init__, (self,) + args, params)
27
28        self.master.title('TkConch')
29        self.localRemoteVar = Tkinter.StringVar()
30        self.localRemoteVar.set('local')
31
32        Tkinter.Label(self, anchor='w', justify='left', text='Hostname').grid(column=1, row=1, sticky='w')
33        self.host = Tkinter.Entry(self)
34        self.host.grid(column=2, columnspan=2, row=1, sticky='nesw')
35
36        Tkinter.Label(self, anchor='w', justify='left', text='Port').grid(column=1, row=2, sticky='w')
37        self.port = Tkinter.Entry(self)
38        self.port.grid(column=2, columnspan=2, row=2, sticky='nesw')
39
40        Tkinter.Label(self, anchor='w', justify='left', text='Username').grid(column=1, row=3, sticky='w')
41        self.user = Tkinter.Entry(self)
42        self.user.grid(column=2, columnspan=2, row=3, sticky='nesw')
43
44        Tkinter.Label(self, anchor='w', justify='left', text='Command').grid(column=1, row=4, sticky='w')
45        self.command = Tkinter.Entry(self)
46        self.command.grid(column=2, columnspan=2, row=4, sticky='nesw')
47
48        Tkinter.Label(self, anchor='w', justify='left', text='Identity').grid(column=1, row=5, sticky='w')
49        self.identity = Tkinter.Entry(self)
50        self.identity.grid(column=2, row=5, sticky='nesw')
51        Tkinter.Button(self, command=self.getIdentityFile, text='Browse').grid(column=3, row=5, sticky='nesw')
52
53        Tkinter.Label(self, text='Port Forwarding').grid(column=1, row=6, sticky='w')
54        self.forwards = Tkinter.Listbox(self, height=0, width=0)
55        self.forwards.grid(column=2, columnspan=2, row=6, sticky='nesw')
56        Tkinter.Button(self, text='Add', command=self.addForward).grid(column=1, row=7)
57        Tkinter.Button(self, text='Remove', command=self.removeForward).grid(column=1, row=8)
58        self.forwardPort = Tkinter.Entry(self)
59        self.forwardPort.grid(column=2, row=7, sticky='nesw')
60        Tkinter.Label(self, text='Port').grid(column=3, row=7, sticky='nesw')
61        self.forwardHost = Tkinter.Entry(self)
62        self.forwardHost.grid(column=2, row=8, sticky='nesw')
63        Tkinter.Label(self, text='Host').grid(column=3, row=8, sticky='nesw')
64        self.localForward = Tkinter.Radiobutton(self, text='Local', variable=self.localRemoteVar, value='local')
65        self.localForward.grid(column=2, row=9)
66        self.remoteForward = Tkinter.Radiobutton(self, text='Remote', variable=self.localRemoteVar, value='remote')
67        self.remoteForward.grid(column=3, row=9)
68
69        Tkinter.Label(self, text='Advanced Options').grid(column=1, columnspan=3, row=10, sticky='nesw')
70
71        Tkinter.Label(self, anchor='w', justify='left', text='Cipher').grid(column=1, row=11, sticky='w')
72        self.cipher = Tkinter.Entry(self, name='cipher')
73        self.cipher.grid(column=2, columnspan=2, row=11, sticky='nesw')
74
75        Tkinter.Label(self, anchor='w', justify='left', text='MAC').grid(column=1, row=12, sticky='w')
76        self.mac = Tkinter.Entry(self, name='mac')
77        self.mac.grid(column=2, columnspan=2, row=12, sticky='nesw')
78
79        Tkinter.Label(self, anchor='w', justify='left', text='Escape Char').grid(column=1, row=13, sticky='w')
80        self.escape = Tkinter.Entry(self, name='escape')
81        self.escape.grid(column=2, columnspan=2, row=13, sticky='nesw')
82        Tkinter.Button(self, text='Connect!', command=self.doConnect).grid(column=1, columnspan=3, row=14, sticky='nesw')
83
84        # Resize behavior(s)
85        self.grid_rowconfigure(6, weight=1, minsize=64)
86        self.grid_columnconfigure(2, weight=1, minsize=2)
87
88        self.master.protocol("WM_DELETE_WINDOW", sys.exit)
89       
90
91    def getIdentityFile(self):
92        r = tkFileDialog.askopenfilename()
93        if r:
94            self.identity.delete(0, Tkinter.END)
95            self.identity.insert(Tkinter.END, r)
96
97    def addForward(self):
98        port = self.forwardPort.get()
99        self.forwardPort.delete(0, Tkinter.END)
100        host = self.forwardHost.get()
101        self.forwardHost.delete(0, Tkinter.END)
102        if self.localRemoteVar.get() == 'local':
103            self.forwards.insert(Tkinter.END, 'L:%s:%s' % (port, host))
104        else:
105            self.forwards.insert(Tkinter.END, 'R:%s:%s' % (port, host))
106
107    def removeForward(self):
108        cur = self.forwards.curselection()
109        if cur:
110            self.forwards.remove(cur[0])
111
112    def doConnect(self):
113        finished = 1
114        options['host'] = self.host.get()
115        options['port'] = self.port.get()
116        options['user'] = self.user.get()
117        options['command'] = self.command.get()
118        cipher = self.cipher.get()
119        mac = self.mac.get()
120        escape = self.escape.get()
121        if cipher:
122            if cipher in SSHClientTransport.supportedCiphers:
123                SSHClientTransport.supportedCiphers = [cipher]
124            else:
125                tkMessageBox.showerror('TkConch', 'Bad cipher.')
126                finished = 0
127
128        if mac:
129            if mac in SSHClientTransport.supportedMACs:
130                SSHClientTransport.supportedMACs = [mac]
131            elif finished:
132                tkMessageBox.showerror('TkConch', 'Bad MAC.')
133                finished = 0
134
135        if escape:
136            if escape == 'none':
137                options['escape'] = None
138            elif escape[0] == '^' and len(escape) == 2:
139                options['escape'] = chr(ord(escape[1])-64)
140            elif len(escape) == 1:
141                options['escape'] = escape
142            elif finished:
143                tkMessageBox.showerror('TkConch', "Bad escape character '%s'." % escape)
144                finished = 0
145
146        if self.identity.get():
147            options.identitys.append(self.identity.get())
148
149        for line in self.forwards.get(0,Tkinter.END):
150            if line[0]=='L':
151                options.opt_localforward(line[2:])
152            else:
153                options.opt_remoteforward(line[2:])
154
155        if '@' in options['host']:
156            options['user'], options['host'] = options['host'].split('@',1)
157
158        if (not options['host'] or not options['user']) and finished:
159            tkMessageBox.showerror('TkConch', 'Missing host or username.')
160            finished = 0
161        if finished:
162            self.master.quit()
163            self.master.destroy()       
164            if options['log']:
165                realout = sys.stdout
166                log.startLogging(sys.stderr)
167                sys.stdout = realout
168            else:
169                log.discardLogs()
170            log.deferr = handleError # HACK
171            if not options.identitys:
172                options.identitys = ['~/.ssh/id_rsa', '~/.ssh/id_dsa']
173            host = options['host']
174            port = int(options['port'] or 22)
175            log.msg((host,port))
176            reactor.connectTCP(host, port, SSHClientFactory())
177            frame.master.deiconify()
178            frame.master.title('%s@%s - TkConch' % (options['user'], options['host']))
179        else:
180            self.focus()
181
182class GeneralOptions(usage.Options):
183    synopsis = """Usage:    tkconch [options] host [command]
184 """
185
186    optParameters = [['user', 'l', None, 'Log in using this user name.'],
187                    ['identity', 'i', '~/.ssh/identity', 'Identity for public key authentication'],
188                    ['escape', 'e', '~', "Set escape character; ``none'' = disable"],
189                    ['cipher', 'c', None, 'Select encryption algorithm.'],
190                    ['macs', 'm', None, 'Specify MAC algorithms for protocol version 2.'],
191                    ['port', 'p', None, 'Connect to this port.  Server must be on the same port.'],
192                    ['localforward', 'L', None, 'listen-port:host:port   Forward local port to remote address'],
193                    ['remoteforward', 'R', None, 'listen-port:host:port   Forward remote port to local address'],
194                    ]
195
196    optFlags = [['tty', 't', 'Tty; allocate a tty even if command is given.'],
197                ['notty', 'T', 'Do not allocate a tty.'],
198                ['version', 'V', 'Display version number only.'],
199                ['compress', 'C', 'Enable compression.'],
200                ['noshell', 'N', 'Do not execute a shell or command.'],
201                ['subsystem', 's', 'Invoke command (mandatory) as SSH2 subsystem.'],
202                ['log', 'v', 'Log to stderr'],
203                ['ansilog', 'a', 'Print the receieved data to stdout']]
204
205    _ciphers = transport.SSHClientTransport.supportedCiphers
206    _macs = transport.SSHClientTransport.supportedMACs
207
208    compData = usage.Completions(
209        mutuallyExclusive=[("tty", "notty")],
210        optActions={
211            "cipher": usage.CompleteList(_ciphers),
212            "macs": usage.CompleteList(_macs),
213            "localforward": usage.Completer(descr="listen-port:host:port"),
214            "remoteforward": usage.Completer(descr="listen-port:host:port")},
215        extraActions=[usage.CompleteUserAtHost(),
216                      usage.Completer(descr="command"),
217                      usage.Completer(descr="argument", repeat=True)]
218        )
219
220    identitys = []
221    localForwards = []
222    remoteForwards = []
223
224    def opt_identity(self, i):
225        self.identitys.append(i)
226
227    def opt_localforward(self, f):
228        localPort, remoteHost, remotePort = f.split(':') # doesn't do v6 yet
229        localPort = int(localPort)
230        remotePort = int(remotePort)
231        self.localForwards.append((localPort, (remoteHost, remotePort)))
232
233    def opt_remoteforward(self, f):
234        remotePort, connHost, connPort = f.split(':') # doesn't do v6 yet
235        remotePort = int(remotePort)
236        connPort = int(connPort)
237        self.remoteForwards.append((remotePort, (connHost, connPort)))
238
239    def opt_compress(self):
240        SSHClientTransport.supportedCompressions[0:1] = ['zlib']
241
242    def parseArgs(self, *args):
243        if args:
244            self['host'] = args[0]
245            self['command'] = ' '.join(args[1:])
246        else:
247            self['host'] = ''
248            self['command'] = ''
249
250# Rest of code in "run"
251options = None
252menu = None
253exitStatus = 0
254frame = None
255
256def deferredAskFrame(question, echo):
257    if frame.callback:
258        raise ValueError("can't ask 2 questions at once!")
259    d = defer.Deferred()
260    resp = []
261    def gotChar(ch, resp=resp):
262        if not ch: return
263        if ch=='\x03': # C-c
264            reactor.stop()
265        if ch=='\r':
266            frame.write('\r\n')
267            stresp = ''.join(resp)
268            del resp
269            frame.callback = None
270            d.callback(stresp)
271            return
272        elif 32 <= ord(ch) < 127:
273            resp.append(ch)
274            if echo:
275                frame.write(ch)
276        elif ord(ch) == 8 and resp: # BS
277            if echo: frame.write('\x08 \x08')
278            resp.pop()
279    frame.callback = gotChar
280    frame.write(question)
281    frame.canvas.focus_force()
282    return d
283
284def run():
285    global menu, options, frame
286    args = sys.argv[1:]
287    if '-l' in args: # cvs is an idiot
288        i = args.index('-l')
289        args = args[i:i+2]+args
290        del args[i+2:i+4]
291    for arg in args[:]:
292        try:
293            i = args.index(arg)
294            if arg[:2] == '-o' and args[i+1][0]!='-':
295                args[i:i+2] = [] # suck on it scp
296        except ValueError:
297            pass
298    root = Tkinter.Tk()
299    root.withdraw()
300    top = Tkinter.Toplevel()
301    menu = TkConchMenu(top)
302    menu.pack(side=Tkinter.TOP, fill=Tkinter.BOTH, expand=1)
303    options = GeneralOptions()
304    try:
305        options.parseOptions(args)
306    except usage.UsageError, u:
307        print 'ERROR: %s' % u
308        options.opt_help()
309        sys.exit(1)
310    for k,v in options.items():
311        if v and hasattr(menu, k):
312            getattr(menu,k).insert(Tkinter.END, v)
313    for (p, (rh, rp)) in options.localForwards:
314        menu.forwards.insert(Tkinter.END, 'L:%s:%s:%s' % (p, rh, rp))
315    options.localForwards = []
316    for (p, (rh, rp)) in options.remoteForwards:
317        menu.forwards.insert(Tkinter.END, 'R:%s:%s:%s' % (p, rh, rp))
318    options.remoteForwards = []
319    frame = tkvt100.VT100Frame(root, callback=None)
320    root.geometry('%dx%d'%(tkvt100.fontWidth*frame.width+3, tkvt100.fontHeight*frame.height+3))
321    frame.pack(side = Tkinter.TOP)
322    tksupport.install(root)
323    root.withdraw()
324    if (options['host'] and options['user']) or '@' in options['host']:
325        menu.doConnect()
326    else:
327        top.mainloop()
328    reactor.run()
329    sys.exit(exitStatus)
330
331def handleError():
332    from twisted.python import failure
333    global exitStatus
334    exitStatus = 2
335    log.err(failure.Failure())
336    reactor.stop()
337    raise
338
339class SSHClientFactory(protocol.ClientFactory):
340    noisy = 1 
341
342    def stopFactory(self):
343        reactor.stop()
344
345    def buildProtocol(self, addr):
346        return SSHClientTransport()
347
348    def clientConnectionFailed(self, connector, reason):
349        tkMessageBox.showwarning('TkConch','Connection Failed, Reason:\n %s: %s' % (reason.type, reason.value))
350
351class SSHClientTransport(transport.SSHClientTransport):
352
353    def receiveError(self, code, desc):
354        global exitStatus
355        exitStatus = 'conch:\tRemote side disconnected with error code %i\nconch:\treason: %s' % (code, desc)
356
357    def sendDisconnect(self, code, reason):
358        global exitStatus
359        exitStatus = 'conch:\tSending disconnect with error code %i\nconch:\treason: %s' % (code, reason)
360        transport.SSHClientTransport.sendDisconnect(self, code, reason)
361
362    def receiveDebug(self, alwaysDisplay, message, lang):
363        global options
364        if alwaysDisplay or options['log']:
365            log.msg('Received Debug Message: %s' % message)
366
367    def verifyHostKey(self, pubKey, fingerprint):
368        #d = defer.Deferred()
369        #d.addCallback(lambda x:defer.succeed(1))
370        #d.callback(2)
371        #return d
372        goodKey = isInKnownHosts(options['host'], pubKey, {'known-hosts': None})
373        if goodKey == 1: # good key
374            return defer.succeed(1)
375        elif goodKey == 2: # AAHHHHH changed
376            return defer.fail(error.ConchError('bad host key'))
377        else:
378            if options['host'] == self.transport.getPeer()[1]:
379                host = options['host']
380                khHost = options['host']
381            else:
382                host = '%s (%s)' % (options['host'], 
383                                    self.transport.getPeer()[1])
384                khHost = '%s,%s' % (options['host'], 
385                                    self.transport.getPeer()[1])
386            keyType = common.getNS(pubKey)[0]
387            ques = """The authenticity of host '%s' can't be established.\r
388%s key fingerprint is %s.""" % (host, 
389                                {'ssh-dss':'DSA', 'ssh-rsa':'RSA'}[keyType], 
390                                fingerprint) 
391            ques+='\r\nAre you sure you want to continue connecting (yes/no)? '
392            return deferredAskFrame(ques, 1).addCallback(self._cbVerifyHostKey, pubKey, khHost, keyType)
393
394    def _cbVerifyHostKey(self, ans, pubKey, khHost, keyType):
395        if ans.lower() not in ('yes', 'no'):
396            return deferredAskFrame("Please type  'yes' or 'no': ",1).addCallback(self._cbVerifyHostKey, pubKey, khHost, keyType)
397        if ans.lower() == 'no':
398            frame.write('Host key verification failed.\r\n')
399            raise error.ConchError('bad host key')
400        try:
401            frame.write("Warning: Permanently added '%s' (%s) to the list of known hosts.\r\n" % (khHost, {'ssh-dss':'DSA', 'ssh-rsa':'RSA'}[keyType]))
402            known_hosts = open(os.path.expanduser('~/.ssh/known_hosts'), 'a')
403            encodedKey = base64.encodestring(pubKey).replace('\n', '')
404            known_hosts.write('\n%s %s %s' % (khHost, keyType, encodedKey))
405            known_hosts.close()
406        except:
407            log.deferr()
408            raise error.ConchError
409
410    def connectionSecure(self):
411        if options['user']:
412            user = options['user']
413        else:
414            user = getpass.getuser()
415        self.requestService(SSHUserAuthClient(user, SSHConnection()))
416
417class SSHUserAuthClient(userauth.SSHUserAuthClient):
418    usedFiles = []
419
420    def getPassword(self, prompt = None):
421        if not prompt:
422            prompt = "%s@%s's password: " % (self.user, options['host'])
423        return deferredAskFrame(prompt,0) 
424
425    def getPublicKey(self):
426        files = [x for x in options.identitys if x not in self.usedFiles]
427        if not files:
428            return None
429        file = files[0]
430        log.msg(file)
431        self.usedFiles.append(file)
432        file = os.path.expanduser(file) 
433        file += '.pub'
434        if not os.path.exists(file):
435            return
436        try:
437            return keys.Key.fromFile(file).blob() 
438        except:
439            return self.getPublicKey() # try again
440   
441    def getPrivateKey(self):
442        file = os.path.expanduser(self.usedFiles[-1])
443        if not os.path.exists(file):
444            return None
445        try:
446            return defer.succeed(keys.Key.fromFile(file).keyObject)
447        except keys.BadKeyError, e:
448            if e.args[0] == 'encrypted key with no password':
449                prompt = "Enter passphrase for key '%s': " % \
450                       self.usedFiles[-1]
451                return deferredAskFrame(prompt, 0).addCallback(self._cbGetPrivateKey, 0)
452    def _cbGetPrivateKey(self, ans, count):
453        file = os.path.expanduser(self.usedFiles[-1])
454        try:
455            return keys.Key.fromFile(file, password = ans).keyObject
456        except keys.BadKeyError:
457            if count == 2:
458                raise
459            prompt = "Enter passphrase for key '%s': " % \
460                   self.usedFiles[-1]
461            return deferredAskFrame(prompt, 0).addCallback(self._cbGetPrivateKey, count+1)
462
463class SSHConnection(connection.SSHConnection):
464    def serviceStarted(self):
465        if not options['noshell']:
466            self.openChannel(SSHSession())
467        if options.localForwards:
468            for localPort, hostport in options.localForwards:
469                reactor.listenTCP(localPort,
470                            forwarding.SSHListenForwardingFactory(self, 
471                                hostport,
472                                forwarding.SSHListenClientForwardingChannel))
473        if options.remoteForwards:
474            for remotePort, hostport in options.remoteForwards:
475                log.msg('asking for remote forwarding for %s:%s' %
476                        (remotePort, hostport))
477                data = forwarding.packGlobal_tcpip_forward(
478                    ('0.0.0.0', remotePort))
479                d = self.sendGlobalRequest('tcpip-forward', data)
480                self.remoteForwards[remotePort] = hostport
481
482class SSHSession(channel.SSHChannel):
483
484    name = 'session'
485   
486    def channelOpen(self, foo):
487        #global globalSession
488        #globalSession = self
489        # turn off local echo
490        self.escapeMode = 1
491        c = session.SSHSessionClient()
492        if options['escape']:
493            c.dataReceived = self.handleInput
494        else:
495            c.dataReceived = self.write
496        c.connectionLost = self.sendEOF
497        frame.callback = c.dataReceived
498        frame.canvas.focus_force()
499        if options['subsystem']:
500            self.conn.sendRequest(self, 'subsystem', \
501                common.NS(options['command']))
502        elif options['command']:
503            if options['tty']:
504                term = os.environ.get('TERM', 'xterm')
505                #winsz = fcntl.ioctl(fd, tty.TIOCGWINSZ, '12345678')
506                winSize = (25,80,0,0) #struct.unpack('4H', winsz)
507                ptyReqData = session.packRequest_pty_req(term, winSize, '')
508                self.conn.sendRequest(self, 'pty-req', ptyReqData)               
509            self.conn.sendRequest(self, 'exec', \
510                common.NS(options['command']))
511        else:
512            if not options['notty']:
513                term = os.environ.get('TERM', 'xterm')
514                #winsz = fcntl.ioctl(fd, tty.TIOCGWINSZ, '12345678')
515                winSize = (25,80,0,0) #struct.unpack('4H', winsz)
516                ptyReqData = session.packRequest_pty_req(term, winSize, '')
517                self.conn.sendRequest(self, 'pty-req', ptyReqData)
518            self.conn.sendRequest(self, 'shell', '')
519        self.conn.transport.transport.setTcpNoDelay(1)
520
521    def handleInput(self, char):
522        #log.msg('handling %s' % repr(char))
523        if char in ('\n', '\r'):
524            self.escapeMode = 1
525            self.write(char)
526        elif self.escapeMode == 1 and char == options['escape']:
527            self.escapeMode = 2
528        elif self.escapeMode == 2:
529            self.escapeMode = 1 # so we can chain escapes together
530            if char == '.': # disconnect
531                log.msg('disconnecting from escape')
532                reactor.stop()
533                return
534            elif char == '\x1a': # ^Z, suspend
535                # following line courtesy of Erwin@freenode
536                os.kill(os.getpid(), signal.SIGSTOP)
537                return
538            elif char == 'R': # rekey connection
539                log.msg('rekeying connection')
540                self.conn.transport.sendKexInit()
541                return
542            self.write('~' + char)
543        else:
544            self.escapeMode = 0
545            self.write(char)
546
547    def dataReceived(self, data):
548        if options['ansilog']:
549            print repr(data)
550        frame.write(data)
551
552    def extReceived(self, t, data):
553        if t==connection.EXTENDED_DATA_STDERR:
554            log.msg('got %s stderr data' % len(data))
555            sys.stderr.write(data)
556            sys.stderr.flush()
557
558    def eofReceived(self):
559        log.msg('got eof')
560        sys.stdin.close()
561
562    def closed(self):
563        log.msg('closed %s' % self)
564        if len(self.conn.channels) == 1: # just us left
565            reactor.stop()
566
567    def request_exit_status(self, data):
568        global exitStatus
569        exitStatus = int(struct.unpack('>L', data)[0])
570        log.msg('exit status: %s' % exitStatus)
571
572    def sendEOF(self):
573        self.conn.sendEOF(self)
574
575if __name__=="__main__":
576    run()
Note: See TracBrowser for help on using the browser.