Statistics
| Branch: | Tag: | Revision:

root / ncclient / transport / ssh.py @ 41e2ed46

History | View | Annotate | Download (10.6 kB)

1 d095a59e Shikhar Bhushan
# Copyright 2009 Shikhar Bhushan
2 d095a59e Shikhar Bhushan
#
3 d095a59e Shikhar Bhushan
# Licensed under the Apache License, Version 2.0 (the "License");
4 d095a59e Shikhar Bhushan
# you may not use this file except in compliance with the License.
5 d095a59e Shikhar Bhushan
# You may obtain a copy of the License at
6 d095a59e Shikhar Bhushan
#
7 d095a59e Shikhar Bhushan
#    http://www.apache.org/licenses/LICENSE-2.0
8 d095a59e Shikhar Bhushan
#
9 d095a59e Shikhar Bhushan
# Unless required by applicable law or agreed to in writing, software
10 d095a59e Shikhar Bhushan
# distributed under the License is distributed on an "AS IS" BASIS,
11 d095a59e Shikhar Bhushan
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 d095a59e Shikhar Bhushan
# See the License for the specific language governing permissions and
13 d095a59e Shikhar Bhushan
# limitations under the License.
14 d095a59e Shikhar Bhushan
15 d095a59e Shikhar Bhushan
import os
16 d095a59e Shikhar Bhushan
import socket
17 d095a59e Shikhar Bhushan
from binascii import hexlify
18 d095a59e Shikhar Bhushan
from cStringIO import StringIO
19 d095a59e Shikhar Bhushan
from select import select
20 d095a59e Shikhar Bhushan
21 d095a59e Shikhar Bhushan
import paramiko
22 d095a59e Shikhar Bhushan
23 c35cebbf Shikhar Bhushan
from errors import AuthenticationError, SessionCloseError, SSHError, SSHUnknownHostError
24 d095a59e Shikhar Bhushan
from session import Session
25 d095a59e Shikhar Bhushan
26 41e2ed46 Shikhar Bhushan
import logging
27 41e2ed46 Shikhar Bhushan
logger = logging.getLogger('ncclient.transport.ssh')
28 41e2ed46 Shikhar Bhushan
29 d095a59e Shikhar Bhushan
BUF_SIZE = 4096
30 d095a59e Shikhar Bhushan
MSG_DELIM = ']]>]]>'
31 d095a59e Shikhar Bhushan
TICK = 0.1
32 d095a59e Shikhar Bhushan
33 d095a59e Shikhar Bhushan
class SSHSession(Session):
34 e91a5349 Shikhar Bhushan
    
35 d095a59e Shikhar Bhushan
    def __init__(self):
36 d095a59e Shikhar Bhushan
        Session.__init__(self)
37 d095a59e Shikhar Bhushan
        self._host_keys = paramiko.HostKeys()
38 4ba5e843 Shikhar Bhushan
        self._system_host_keys = paramiko.HostKeys()
39 d095a59e Shikhar Bhushan
        self._transport = None
40 d095a59e Shikhar Bhushan
        self._connected = False
41 d095a59e Shikhar Bhushan
        self._channel = None
42 1d540e60 Shikhar Bhushan
        self._expecting_close = False
43 d095a59e Shikhar Bhushan
        self._buffer = StringIO() # for incoming data
44 d095a59e Shikhar Bhushan
        # parsing-related, see _parse()
45 d095a59e Shikhar Bhushan
        self._parsing_state = 0 
46 d095a59e Shikhar Bhushan
        self._parsing_pos = 0
47 d095a59e Shikhar Bhushan
    
48 d095a59e Shikhar Bhushan
    def _parse(self):
49 d095a59e Shikhar Bhushan
        '''Messages ae delimited by MSG_DELIM. The buffer could have grown by a
50 d095a59e Shikhar Bhushan
        maximum of BUF_SIZE bytes everytime this method is called. Retains state
51 d095a59e Shikhar Bhushan
        across method calls and if a byte has been read it will not be considered
52 d095a59e Shikhar Bhushan
        again.
53 d095a59e Shikhar Bhushan
        '''
54 d095a59e Shikhar Bhushan
        delim = MSG_DELIM
55 d095a59e Shikhar Bhushan
        n = len(delim) - 1
56 d095a59e Shikhar Bhushan
        expect = self._parsing_state
57 d095a59e Shikhar Bhushan
        buf = self._buffer
58 d095a59e Shikhar Bhushan
        buf.seek(self._parsing_pos)
59 d095a59e Shikhar Bhushan
        while True:
60 d095a59e Shikhar Bhushan
            x = buf.read(1)
61 d095a59e Shikhar Bhushan
            if not x: # done reading
62 d095a59e Shikhar Bhushan
                break
63 d095a59e Shikhar Bhushan
            elif x == delim[expect]: # what we expected
64 d095a59e Shikhar Bhushan
                expect += 1 # expect the next delim char
65 d095a59e Shikhar Bhushan
            else:
66 d095a59e Shikhar Bhushan
                continue
67 d095a59e Shikhar Bhushan
            # loop till last delim char expected, break if other char encountered
68 d095a59e Shikhar Bhushan
            for i in range(expect, n):
69 d095a59e Shikhar Bhushan
                x = buf.read(1)
70 d095a59e Shikhar Bhushan
                if not x: # done reading
71 d095a59e Shikhar Bhushan
                    break
72 d095a59e Shikhar Bhushan
                if x == delim[expect]: # what we expected
73 d095a59e Shikhar Bhushan
                    expect += 1 # expect the next delim char
74 d095a59e Shikhar Bhushan
                else:
75 d095a59e Shikhar Bhushan
                    expect = 0 # reset
76 d095a59e Shikhar Bhushan
                    break
77 d095a59e Shikhar Bhushan
            else: # if we didn't break out of the loop, full delim was parsed
78 d095a59e Shikhar Bhushan
                msg_till = buf.tell() - n
79 d095a59e Shikhar Bhushan
                buf.seek(0)
80 41e2ed46 Shikhar Bhushan
                logger.debug('parsed new message')
81 41e2ed46 Shikhar Bhushan
                self._dispatch_message(buf.read(msg_till).strip())
82 d095a59e Shikhar Bhushan
                buf.seek(n+1, os.SEEK_CUR)
83 d095a59e Shikhar Bhushan
                rest = buf.read()
84 d095a59e Shikhar Bhushan
                buf = StringIO()
85 d095a59e Shikhar Bhushan
                buf.write(rest)
86 d095a59e Shikhar Bhushan
                buf.seek(0)
87 4ba5e843 Shikhar Bhushan
                expect = 0
88 d095a59e Shikhar Bhushan
        self._buffer = buf
89 d095a59e Shikhar Bhushan
        self._parsing_state = expect
90 d095a59e Shikhar Bhushan
        self._parsing_pos = self._buffer.tell()
91 d095a59e Shikhar Bhushan
    
92 1d540e60 Shikhar Bhushan
    def expect_close(self):
93 1d540e60 Shikhar Bhushan
        self._expecting_close = True
94 1d540e60 Shikhar Bhushan
    
95 d095a59e Shikhar Bhushan
    def load_system_host_keys(self, filename=None):
96 d095a59e Shikhar Bhushan
        if filename is None:
97 d095a59e Shikhar Bhushan
            filename = os.path.expanduser('~/.ssh/known_hosts')
98 d095a59e Shikhar Bhushan
            try:
99 d095a59e Shikhar Bhushan
                self._system_host_keys.load(filename)
100 d095a59e Shikhar Bhushan
            except IOError:
101 4ba5e843 Shikhar Bhushan
                # for windows
102 4ba5e843 Shikhar Bhushan
                filename = os.path.expanduser('~/ssh/known_hosts')
103 4ba5e843 Shikhar Bhushan
                try:
104 4ba5e843 Shikhar Bhushan
                    self._system_host_keys.load(filename)
105 4ba5e843 Shikhar Bhushan
                except IOError:
106 4ba5e843 Shikhar Bhushan
                    pass
107 d095a59e Shikhar Bhushan
            return
108 d095a59e Shikhar Bhushan
        self._system_host_keys.load(filename)
109 d095a59e Shikhar Bhushan
    
110 d095a59e Shikhar Bhushan
    def load_host_keys(self, filename):
111 d095a59e Shikhar Bhushan
        self._host_keys.load(filename)
112 d095a59e Shikhar Bhushan
113 d095a59e Shikhar Bhushan
    def add_host_key(self, key):
114 d095a59e Shikhar Bhushan
        self._host_keys.add(key)
115 d095a59e Shikhar Bhushan
    
116 d095a59e Shikhar Bhushan
    def save_host_keys(self, filename):
117 d095a59e Shikhar Bhushan
        f = open(filename, 'w')
118 d095a59e Shikhar Bhushan
        for hostname, keys in self._host_keys.iteritems():
119 d095a59e Shikhar Bhushan
            for keytype, key in keys.iteritems():
120 d095a59e Shikhar Bhushan
                f.write('%s %s %s\n' % (hostname, keytype, key.get_base64()))
121 d095a59e Shikhar Bhushan
        f.close()    
122 d095a59e Shikhar Bhushan
    
123 d095a59e Shikhar Bhushan
    def close(self):
124 d095a59e Shikhar Bhushan
        if self._transport.is_active():
125 d095a59e Shikhar Bhushan
            self._transport.close()
126 d095a59e Shikhar Bhushan
        self._connected = False
127 d095a59e Shikhar Bhushan
    
128 d095a59e Shikhar Bhushan
    def connect(self, hostname, port=830, timeout=None,
129 d095a59e Shikhar Bhushan
                unknown_host_cb=None, username=None, password=None,
130 d095a59e Shikhar Bhushan
                key_filename=None, allow_agent=True, look_for_keys=True):
131 d095a59e Shikhar Bhushan
        
132 d095a59e Shikhar Bhushan
        assert(username is not None)
133 d095a59e Shikhar Bhushan
        
134 d095a59e Shikhar Bhushan
        for (family, socktype, proto, canonname, sockaddr) in \
135 d095a59e Shikhar Bhushan
        socket.getaddrinfo(hostname, port):
136 d095a59e Shikhar Bhushan
            if socktype==socket.SOCK_STREAM:
137 d095a59e Shikhar Bhushan
                af = family
138 d095a59e Shikhar Bhushan
                addr = sockaddr
139 d095a59e Shikhar Bhushan
                break
140 d095a59e Shikhar Bhushan
        else:
141 d095a59e Shikhar Bhushan
            raise SSHError('No suitable address family for %s' % hostname)
142 d095a59e Shikhar Bhushan
        sock = socket.socket(af, socket.SOCK_STREAM)
143 d095a59e Shikhar Bhushan
        sock.settimeout(timeout)
144 d095a59e Shikhar Bhushan
        sock.connect(addr)
145 d095a59e Shikhar Bhushan
        t = self._transport = paramiko.Transport(sock)
146 d095a59e Shikhar Bhushan
        t.set_log_channel(logger.name)
147 d095a59e Shikhar Bhushan
        
148 d095a59e Shikhar Bhushan
        try:
149 d095a59e Shikhar Bhushan
            t.start_client()
150 d095a59e Shikhar Bhushan
        except paramiko.SSHException:
151 d095a59e Shikhar Bhushan
            raise SSHError('Negotiation failed')
152 d095a59e Shikhar Bhushan
        
153 d095a59e Shikhar Bhushan
        # host key verification
154 d095a59e Shikhar Bhushan
        server_key = t.get_remote_server_key()
155 d095a59e Shikhar Bhushan
        known_host = self._host_keys.check(hostname, server_key) or \
156 d095a59e Shikhar Bhushan
                        self._system_host_keys.check(hostname, server_key)
157 d095a59e Shikhar Bhushan
        
158 d095a59e Shikhar Bhushan
        if unknown_host_cb is None:
159 d095a59e Shikhar Bhushan
            unknown_host_cb = lambda *args: False
160 d095a59e Shikhar Bhushan
        if not known_host and not unknown_host_cb(hostname, server_key):
161 d095a59e Shikhar Bhushan
                raise SSHUnknownHostError(hostname, server_key)
162 d095a59e Shikhar Bhushan
        
163 d095a59e Shikhar Bhushan
        if key_filename is None:
164 d095a59e Shikhar Bhushan
            key_filenames = []
165 d095a59e Shikhar Bhushan
        elif isinstance(key_filename, basestring):
166 d095a59e Shikhar Bhushan
            key_filenames = [ key_filename ]
167 d095a59e Shikhar Bhushan
        else:
168 d095a59e Shikhar Bhushan
            key_filenames = key_filename
169 d095a59e Shikhar Bhushan
        
170 d095a59e Shikhar Bhushan
        self._auth(username, password, key_filenames, allow_agent, look_for_keys)
171 d095a59e Shikhar Bhushan
        
172 d095a59e Shikhar Bhushan
        self._connected = True # there was no error authenticating
173 d095a59e Shikhar Bhushan
        
174 d095a59e Shikhar Bhushan
        c = self._channel = self._transport.open_session()
175 d095a59e Shikhar Bhushan
        c.set_name('netconf')
176 41e2ed46 Shikhar Bhushan
        c.invoke_subsystem('netconf')
177 d095a59e Shikhar Bhushan
        
178 d095a59e Shikhar Bhushan
        self._post_connect()
179 d095a59e Shikhar Bhushan
    
180 d095a59e Shikhar Bhushan
    # on the lines of paramiko.SSHClient._auth()
181 d095a59e Shikhar Bhushan
    def _auth(self, username, password, key_filenames, allow_agent,
182 d095a59e Shikhar Bhushan
              look_for_keys):
183 d095a59e Shikhar Bhushan
        saved_exception = None
184 d095a59e Shikhar Bhushan
        
185 d095a59e Shikhar Bhushan
        for key_filename in key_filenames:
186 d095a59e Shikhar Bhushan
            for cls in (paramiko.RSAKey, paramiko.DSSKey):
187 d095a59e Shikhar Bhushan
                try:
188 d095a59e Shikhar Bhushan
                    key = cls.from_private_key_file(key_filename, password)
189 d095a59e Shikhar Bhushan
                    logger.debug('Trying key %s from %s' %
190 d095a59e Shikhar Bhushan
                              (hexlify(key.get_fingerprint()), key_filename))
191 d095a59e Shikhar Bhushan
                    self._transport.auth_publickey(username, key)
192 d095a59e Shikhar Bhushan
                    return
193 d095a59e Shikhar Bhushan
                except Exception as e:
194 d095a59e Shikhar Bhushan
                    saved_exception = e
195 d095a59e Shikhar Bhushan
                    logger.debug(e)
196 d095a59e Shikhar Bhushan
        
197 d095a59e Shikhar Bhushan
        if allow_agent:
198 d095a59e Shikhar Bhushan
            for key in paramiko.Agent().get_keys():
199 d095a59e Shikhar Bhushan
                try:
200 d095a59e Shikhar Bhushan
                    logger.debug('Trying SSH agent key %s' %
201 d095a59e Shikhar Bhushan
                                 hexlify(key.get_fingerprint()))
202 d095a59e Shikhar Bhushan
                    self._transport.auth_publickey(username, key)
203 d095a59e Shikhar Bhushan
                    return
204 d095a59e Shikhar Bhushan
                except Exception as e:
205 d095a59e Shikhar Bhushan
                    saved_exception = e
206 d095a59e Shikhar Bhushan
                    logger.debug(e)
207 d095a59e Shikhar Bhushan
        
208 d095a59e Shikhar Bhushan
        keyfiles = []
209 d095a59e Shikhar Bhushan
        if look_for_keys:
210 d095a59e Shikhar Bhushan
            rsa_key = os.path.expanduser('~/.ssh/id_rsa')
211 d095a59e Shikhar Bhushan
            dsa_key = os.path.expanduser('~/.ssh/id_dsa')
212 d095a59e Shikhar Bhushan
            if os.path.isfile(rsa_key):
213 d095a59e Shikhar Bhushan
                keyfiles.append((paramiko.RSAKey, rsa_key))
214 d095a59e Shikhar Bhushan
            if os.path.isfile(dsa_key):
215 d095a59e Shikhar Bhushan
                keyfiles.append((paramiko.DSSKey, dsa_key))
216 d095a59e Shikhar Bhushan
            # look in ~/ssh/ for windows users:
217 d095a59e Shikhar Bhushan
            rsa_key = os.path.expanduser('~/ssh/id_rsa')
218 d095a59e Shikhar Bhushan
            dsa_key = os.path.expanduser('~/ssh/id_dsa')
219 d095a59e Shikhar Bhushan
            if os.path.isfile(rsa_key):
220 d095a59e Shikhar Bhushan
                keyfiles.append((paramiko.RSAKey, rsa_key))
221 d095a59e Shikhar Bhushan
            if os.path.isfile(dsa_key):
222 d095a59e Shikhar Bhushan
                keyfiles.append((paramiko.DSSKey, dsa_key))
223 d095a59e Shikhar Bhushan
        
224 d095a59e Shikhar Bhushan
        for cls, filename in keyfiles:
225 d095a59e Shikhar Bhushan
            try:
226 d095a59e Shikhar Bhushan
                key = cls.from_private_key_file(filename, password)
227 d095a59e Shikhar Bhushan
                logger.debug('Trying discovered key %s in %s' %
228 d095a59e Shikhar Bhushan
                          (hexlify(key.get_fingerprint()), filename))
229 d095a59e Shikhar Bhushan
                self._transport.auth_publickey(username, key)
230 d095a59e Shikhar Bhushan
                return
231 d095a59e Shikhar Bhushan
            except Exception as e:
232 d095a59e Shikhar Bhushan
                saved_exception = e
233 d095a59e Shikhar Bhushan
                logger.debug(e)
234 d095a59e Shikhar Bhushan
        
235 d095a59e Shikhar Bhushan
        if password is not None:
236 d095a59e Shikhar Bhushan
            try:
237 d095a59e Shikhar Bhushan
                self._transport.auth_password(username, password)
238 d095a59e Shikhar Bhushan
                return
239 d095a59e Shikhar Bhushan
            except Exception as e:
240 d095a59e Shikhar Bhushan
                saved_exception = e
241 d095a59e Shikhar Bhushan
                logger.debug(e)
242 d095a59e Shikhar Bhushan
        
243 d095a59e Shikhar Bhushan
        if saved_exception is not None:
244 4ba5e843 Shikhar Bhushan
            raise SSHAuthenticationError(repr(saved_exception))
245 d095a59e Shikhar Bhushan
        
246 4ba5e843 Shikhar Bhushan
        raise SSHAuthenticationError('No authentication methods available')
247 d095a59e Shikhar Bhushan
    
248 d095a59e Shikhar Bhushan
    def run(self):
249 d095a59e Shikhar Bhushan
        chan = self._channel
250 d095a59e Shikhar Bhushan
        chan.setblocking(0)
251 d095a59e Shikhar Bhushan
        q = self._q
252 d095a59e Shikhar Bhushan
        try:
253 d095a59e Shikhar Bhushan
            while True:
254 d095a59e Shikhar Bhushan
                # select on a paramiko ssh channel object does not ever return
255 d095a59e Shikhar Bhushan
                # it in the writable list, so it channel's don't exactly emulate 
256 d095a59e Shikhar Bhushan
                # the socket api
257 d095a59e Shikhar Bhushan
                r, w, e = select([chan], [], [], TICK)
258 d095a59e Shikhar Bhushan
                # will wakeup evey TICK seconds to check if something
259 d095a59e Shikhar Bhushan
                # to send, more if something to read (due to select returning
260 d095a59e Shikhar Bhushan
                # chan in readable list)
261 d095a59e Shikhar Bhushan
                if r:
262 d095a59e Shikhar Bhushan
                    data = chan.recv(BUF_SIZE)
263 d095a59e Shikhar Bhushan
                    if data:
264 d095a59e Shikhar Bhushan
                        self._buffer.write(data)
265 d095a59e Shikhar Bhushan
                        self._parse()
266 d095a59e Shikhar Bhushan
                    else:
267 94265508 Shikhar Bhushan
                        raise SessionCloseError(self._buffer.getvalue())
268 d095a59e Shikhar Bhushan
                if not q.empty() and chan.send_ready():
269 41e2ed46 Shikhar Bhushan
                    logger.debug('sending message')
270 d095a59e Shikhar Bhushan
                    data = q.get() + MSG_DELIM
271 d095a59e Shikhar Bhushan
                    while data:
272 d095a59e Shikhar Bhushan
                        n = chan.send(data)
273 d095a59e Shikhar Bhushan
                        if n <= 0:
274 94265508 Shikhar Bhushan
                            raise SessionCloseError(self._buffer.getvalue(), data)
275 d095a59e Shikhar Bhushan
                        data = data[n:]
276 d095a59e Shikhar Bhushan
        except Exception as e:
277 41e2ed46 Shikhar Bhushan
            logger.debug('broke out of main loop')
278 1d540e60 Shikhar Bhushan
            self.close()
279 1d540e60 Shikhar Bhushan
            if not (isinstance(e, SessionCloseError) and self._expecting_close):
280 1d540e60 Shikhar Bhushan
                self._dispatch_error(e)
281 d095a59e Shikhar Bhushan
    
282 d095a59e Shikhar Bhushan
    @property
283 d095a59e Shikhar Bhushan
    def transport(self):
284 e91a5349 Shikhar Bhushan
        '''Get underlying paramiko transport object; this is provided so methods
285 e91a5349 Shikhar Bhushan
        like set_keepalive can be called on it. See paramiko.Transport
286 e91a5349 Shikhar Bhushan
        documentation for details.
287 d095a59e Shikhar Bhushan
        '''
288 d095a59e Shikhar Bhushan
        return self._transport