blob: 6515e4911f74292aca98d8de75c5a46a6c8b5cb1 [file] [log] [blame]
Zsolt Haraszti3bfff662016-12-14 23:41:49 -08001#
2# Copyright 2016 the original author or authors.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8# http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15#
16import os
17import rlcompleter
18from pprint import pprint
19
20import structlog
21from twisted.conch import manhole_ssh
22from twisted.conch.manhole import ColoredManhole
23from twisted.conch.ssh import keys
24from twisted.cred.checkers import InMemoryUsernamePasswordDatabaseDontUse
25from twisted.cred.portal import Portal
26from twisted.internet import reactor
27
28log = structlog.get_logger()
29
30
31MANHOLE_SERVER_RSA_PRIVATE = './manhole_rsa_key'
32MANHOLE_SERVER_RSA_PUBLIC = './manhole_rsa_key.pub'
33
34
35def get_rsa_keys():
36 if not (os.path.exists(MANHOLE_SERVER_RSA_PUBLIC) and os.path.exists(MANHOLE_SERVER_RSA_PRIVATE)):
37 # generate a RSA keypair
38
39 log.info('generate-rsa-keypair')
40 from Crypto.PublicKey import RSA
41 rsa_key = RSA.generate(1024)
42 public_key_str = rsa_key.publickey().exportKey(format='OpenSSH')
43 private_key_str = rsa_key.exportKey()
44
45 # save keys for next time
46 file(MANHOLE_SERVER_RSA_PUBLIC, 'w+b').write(public_key_str)
47 file(MANHOLE_SERVER_RSA_PRIVATE, 'w+b').write(private_key_str)
48 log.debug('saved-rsa-keypair', public=MANHOLE_SERVER_RSA_PUBLIC,
49 private=MANHOLE_SERVER_RSA_PRIVATE)
50 else:
51 public_key_str = file(MANHOLE_SERVER_RSA_PUBLIC).read()
52 private_key_str = file(MANHOLE_SERVER_RSA_PRIVATE).read()
53 return public_key_str, private_key_str
54
55
56class ManholeWithCompleter(ColoredManhole):
57
58 def __init__(self, namespace):
59 namespace['manhole'] = self
60 super(ManholeWithCompleter, self).__init__(namespace)
61 self.last_tab = None
62 self.completer = rlcompleter.Completer(self.namespace)
63
64 def handle_TAB(self):
65 if self.last_tab != self.lineBuffer:
66 self.last_tab = self.lineBuffer
67 return
68
69 buffer = ''.join(self.lineBuffer)
70 completions = []
71 maxlen = 3
72 for c in xrange(1000):
73 candidate = self.completer.complete(buffer, c)
74 if not candidate:
75 break
76
77 if len(candidate) > maxlen:
78 maxlen = len(candidate)
79
80 completions.append(candidate)
81
82 if len(completions) == 1:
83 rest = completions[0][len(buffer):]
84 self.terminal.write(rest)
85 self.lineBufferIndex += len(rest)
86 self.lineBuffer.extend(rest)
87
88 elif len(completions):
89 maxlen += 3
90 numcols = self.width / maxlen
91 self.terminal.nextLine()
92 for idx, candidate in enumerate(completions):
93 self.terminal.write('%%-%ss' % maxlen % candidate)
94 if not ((idx + 1) % numcols):
95 self.terminal.nextLine()
96 self.terminal.nextLine()
97 self.drawInputLine()
98
99
100class Manhole(object):
101
102 def __init__(self, port, pws, **kw):
103 kw.update(globals())
104 kw['pp'] = pprint
105
106 realm = manhole_ssh.TerminalRealm()
107 manhole = ManholeWithCompleter(kw)
108
109 def windowChanged(_, win_size):
110 manhole.terminalSize(*reversed(win_size[:2]))
111
112 realm.sessionFactory.windowChanged = windowChanged
113 realm.chainedProtocolFactory.protocolFactory = lambda _: manhole
114 portal = Portal(realm)
115 portal.registerChecker(InMemoryUsernamePasswordDatabaseDontUse(**pws))
116 factory = manhole_ssh.ConchFactory(portal)
117 public_key_str, private_key_str = get_rsa_keys()
118 factory.publicKeys = {
119 'ssh-rsa': keys.Key.fromString(public_key_str)
120 }
121 factory.privateKeys = {
122 'ssh-rsa': keys.Key.fromString(private_key_str)
123 }
124 reactor.listenTCP(port, factory, interface='localhost')
125
126
127if __name__ == '__main__':
128 Manhole(12222, dict(admin='admin'))
129 reactor.run()