git-svn-id: http://ncclient.googlecode.com/svn/trunk@22 6dbcf712-26ac-11de-a2f3-13738...
[ncclient] / src / ssh.py
1 # Copyright 2009 Shikhar Bhushan
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #    http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14
15 import logging
16 import paramiko
17
18 from select import select as select
19
20 from session import Session, SessionError
21
22 logger = logging.getLogger('ncclient.ssh')
23
24 class SSHError(SessionError): pass
25
26 class SSHSession(Session):
27     
28     BUF_SIZE = 4096
29     MSG_DELIM = ']]>>]]>'
30     MSG_DELIM_LEN = len(MSG_DELIM)
31     
32     def __init__(self, capabilities, load_known_hosts=True,
33                  missing_host_key_policy=paramiko.RejectPolicy):
34         Session.__init__(self, capabilities)
35         self._inBuf = ''
36         self._outBuf = ''
37         self._client = SSHClient()
38         if load_known_hosts:
39             self._client.load_system_host_keys()
40         self._client.set_missing_host_key_policy(missing_host_key_policy)
41     
42     def load_host_keys(self, filename):
43         self._client.load_host_keys(filename)
44     
45     def set_missing_host_key_policy(self, policy):
46         self._client.set_missing_host_key_policy(policy)
47     
48     def connect(self, hostname, port=830, username=None, password=None,
49                 key_filename=None, timeout=None, allow_agent=True,
50                 look_for_keys=True):
51         self._client.connect(hostname, port=port, username=username,
52                              password=password, key_filename=key_filename,
53                              timeout=timeout, allow_agent=allow_agent,
54                              look_for_keys=look_for_keys)
55         transport = self._client.get_transport()
56         self._channel = transport.open_session()
57         self._channel.invoke_subsystem('netconf')
58         self.connected = True
59         self._greet()
60         self.start()
61
62     def _close(self):
63         self._channel.shutdown(2)
64     
65     def run(self):
66         sock = self._channel
67         sock.setblocking(0)
68         q = self._q
69         while True:
70             (r, w, e) = select([sock], [sock], [], 60)
71             if w:
72                 if not q.empty():
73                    self._outBuffer += ( q.get() + MSG_DELIM )
74                 if self._outBuffer:
75                     n = sock.send(self._outBuffer)
76                     self._outBuffer = self._outBuffer[n:]
77             if r:
78                 data = sock.recv(BUF_SIZE)
79                 if data:
80                     self._inBuf += data
81                     (before, _, after) = self._inBuf.partition(MSG_DELIM)
82                     if after:
83                         self.dispatch('reply', before)
84                         self._inBuf = after
85                 else:
86                     self.dispatch('error', self._inBuf)
87
88
89 class CallbackPolicy(paramiko.MissingHostKeyPolicy):
90     
91     def __init__(self, cb):
92         self._cb = cb
93     
94     def missing_host_key(self, client, hostname, key):
95         if not self._cb(hostname, key):
96             raise SSHError