1# Copyright (C) 2011  John Rochester <john@jrochester.org>
2#
3# This file is part of ssh.
4#
5# 'ssh' is free software; you can redistribute it and/or modify it under the
6# terms of the GNU Lesser General Public License as published by the Free
7# Software Foundation; either version 2.1 of the License, or (at your option)
8# any later version.
9#
10# 'ssh' is distrubuted in the hope that it will be useful, but WITHOUT ANY
11# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
12# A PARTICULAR PURPOSE.  See the GNU Lesser General Public License for more
13# details.
14#
15# You should have received a copy of the GNU Lesser General Public License
16# along with 'ssh'; if not, write to the Free Software Foundation, Inc.,
17# 51 Franklin Street, Suite 500, Boston, MA  02110-1335  USA.
18
19"""
20SSH Agent interface for Unix clients.
21"""
22
23import os
24import socket
25import struct
26import sys
27import threading
28import time
29import tempfile
30import stat
31from select import select
32
33from ssh.ssh_exception import SSHException
34from ssh.message import Message
35from ssh.pkey import PKey
36from ssh.channel import Channel
37from ssh.common import io_sleep
38from ssh.util import retry_on_signal
39
40SSH2_AGENTC_REQUEST_IDENTITIES, SSH2_AGENT_IDENTITIES_ANSWER, \
41    SSH2_AGENTC_SIGN_REQUEST, SSH2_AGENT_SIGN_RESPONSE = range(11, 15)
42
43class AgentSSH(object):
44    """
45    Client interface for using private keys from an SSH agent running on the
46    local machine.  If an SSH agent is running, this class can be used to
47    connect to it and retreive L{PKey} objects which can be used when
48    attempting to authenticate to remote SSH servers.
49
50    Because the SSH agent protocol uses environment variables and unix-domain
51    sockets, this probably doesn't work on Windows.  It does work on most
52    posix platforms though (Linux and MacOS X, for example).
53    """
54    def __init__(self):
55        self._conn = None
56        self._keys = ()
57
58    def get_keys(self):
59        """
60        Return the list of keys available through the SSH agent, if any.  If
61        no SSH agent was running (or it couldn't be contacted), an empty list
62        will be returned.
63
64        @return: a list of keys available on the SSH agent
65        @rtype: tuple of L{AgentKey}
66        """
67        return self._keys
68
69    def _connect(self, conn):
70        self._conn = conn
71        ptype, result = self._send_message(chr(SSH2_AGENTC_REQUEST_IDENTITIES))
72        if ptype != SSH2_AGENT_IDENTITIES_ANSWER:
73            raise SSHException('could not get keys from ssh-agent')
74        keys = []
75        for i in range(result.get_int()):
76            keys.append(AgentKey(self, result.get_string()))
77            result.get_string()
78        self._keys = tuple(keys)
79
80    def _close(self):
81        #self._conn.close()
82        self._conn = None
83        self._keys = ()
84
85    def _send_message(self, msg):
86        msg = str(msg)
87        self._conn.send(struct.pack('>I', len(msg)) + msg)
88        l = self._read_all(4)
89        msg = Message(self._read_all(struct.unpack('>I', l)[0]))
90        return ord(msg.get_byte()), msg
91
92    def _read_all(self, wanted):
93        result = self._conn.recv(wanted)
94        while len(result) < wanted:
95            if len(result) == 0:
96                raise SSHException('lost ssh-agent')
97            extra = self._conn.recv(wanted - len(result))
98            if len(extra) == 0:
99                raise SSHException('lost ssh-agent')
100            result += extra
101        return result
102
103class AgentProxyThread(threading.Thread):
104    """ Class in charge of communication between two chan """
105    def __init__(self, agent):
106        threading.Thread.__init__(self, target=self.run)
107        self._agent = agent
108        self._exit = False
109
110    def run(self):
111        try:
112            (r,addr) = self.get_connection()
113            self.__inr = r
114            self.__addr = addr
115            self._agent.connect()
116            self._communicate()
117        except:
118            #XXX Not sure what to do here ... raise or pass ?
119            raise
120
121    def _communicate(self):
122        import fcntl
123        oldflags = fcntl.fcntl(self.__inr, fcntl.F_GETFL)
124        fcntl.fcntl(self.__inr, fcntl.F_SETFL, oldflags | os.O_NONBLOCK)
125        while not self._exit:
126            events = select([self._agent._conn, self.__inr], [], [], 0.5)
127            for fd in events[0]:
128                if self._agent._conn == fd:
129                    data = self._agent._conn.recv(512)
130                    if len(data) != 0:
131                        self.__inr.send(data)
132                    else:
133                        break
134                elif self.__inr == fd:
135                    data = self.__inr.recv(512)
136                    if len(data) != 0:
137                        self._agent._conn.send(data)
138                    else:
139                        break
140            time.sleep(io_sleep)
141
142class AgentLocalProxy(AgentProxyThread):
143    """
144    Class to be used when wanting to ask a local SSH Agent being
145    asked from a remote fake agent (so use a unix socket for ex.)
146    """
147    def __init__(self, agent):
148        AgentProxyThread.__init__(self, agent)
149
150    def get_connection(self):
151        """ Return a pair of socket object and string address
152        May Block !
153        """
154        conn = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
155        try:
156            conn.bind(self._agent._get_filename())
157            conn.listen(1)
158            (r,addr) = conn.accept()
159            return (r, addr)
160        except:
161            raise
162        return None
163
164class AgentRemoteProxy(AgentProxyThread):
165    """
166    Class to be used when wanting to ask a remote SSH Agent
167    """
168    def __init__(self, agent, chan):
169        AgentProxyThread.__init__(self, agent)
170        self.__chan = chan
171
172    def get_connection(self):
173        """
174        Class to be used when wanting to ask a local SSH Agent being
175        asked from a remote fake agent (so use a unix socket for ex.)
176        """
177        return (self.__chan, None)
178
179class AgentClientProxy(object):
180    """
181    Class proxying request as a client:
182       -> client ask for a request_forward_agent()
183       -> server creates a proxy and a fake SSH Agent
184       -> server ask for establishing a connection when needed,
185       calling the forward_agent_handler at client side.
186       -> the forward_agent_handler launch a thread for connecting
187       the remote fake agent and the local agent
188       -> Communication occurs ...
189    """
190    def __init__(self, chanRemote):
191        self._conn = None
192        self.__chanR = chanRemote
193        self.thread = AgentRemoteProxy(self, chanRemote)
194        self.thread.start()
195
196    def __del__(self):
197        self.close()
198
199    def connect(self):
200        """
201        Method automatically called by the run() method of the AgentProxyThread
202        """
203        if ('SSH_AUTH_SOCK' in os.environ) and (sys.platform != 'win32'):
204            conn = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
205            try:
206                retry_on_signal(lambda: conn.connect(os.environ['SSH_AUTH_SOCK']))
207            except:
208                # probably a dangling env var: the ssh agent is gone
209                return
210        elif sys.platform == 'win32':
211            import win_pageant
212            if win_pageant.can_talk_to_agent():
213                conn = win_pageant.PageantConnection()
214            else:
215                return
216        else:
217            # no agent support
218            return
219        self._conn = conn
220
221    def close(self):
222        """
223        Close the current connection and terminate the agent
224        Should be called manually
225        """
226        if hasattr(self, "thread"):
227            self.thread._exit = True
228            self.thread.join(1000)
229        if self._conn is not None:
230            self._conn.close()
231
232class AgentServerProxy(AgentSSH):
233    """
234    @param t : transport used for the Forward for SSH Agent communication
235
236    @raise SSHException: mostly if we lost the agent
237    """
238    def __init__(self, t):
239        AgentSSH.__init__(self)
240        self.__t = t
241        self._dir = tempfile.mkdtemp('sshproxy')
242        os.chmod(self._dir, stat.S_IRWXU)
243        self._file = self._dir + '/sshproxy.ssh'
244        self.thread = AgentLocalProxy(self)
245        self.thread.start()
246
247    def __del__(self):
248        self.close()
249
250    def connect(self):
251         conn_sock = self.__t.open_forward_agent_channel()
252         if conn_sock is None:
253             raise SSHException('lost ssh-agent')
254         conn_sock.set_name('auth-agent')
255         self._connect(conn_sock)
256
257    def close(self):
258        """
259        Terminate the agent, clean the files, close connections
260        Should be called manually
261        """
262        os.remove(self._file)
263        os.rmdir(self._dir)
264        self.thread._exit = True
265        self.thread.join(1000)
266        self._close()
267
268    def get_env(self):
269        """
270        Helper for the environnement under unix
271
272        @return: the SSH_AUTH_SOCK Environnement variables
273        @rtype: dict
274        """
275        env = {}
276        env['SSH_AUTH_SOCK'] = self._get_filename()
277        return env
278
279    def _get_filename(self):
280        return self._file
281
282class AgentRequestHandler(object):
283    def __init__(self, chanClient):
284        self._conn = None
285        self.__chanC = chanClient
286        chanClient.request_forward_agent(self._forward_agent_handler)
287        self.__clientProxys = []
288
289    def _forward_agent_handler(self, chanRemote):
290        self.__clientProxys.append(AgentClientProxy(chanRemote))
291
292    def __del__(self):
293        self.close()
294
295    def close(self):
296        for p in self.__clientProxys:
297            p.close()
298
299class Agent(AgentSSH):
300    """
301    Client interface for using private keys from an SSH agent running on the
302    local machine.  If an SSH agent is running, this class can be used to
303    connect to it and retreive L{PKey} objects which can be used when
304    attempting to authenticate to remote SSH servers.
305
306    Because the SSH agent protocol uses environment variables and unix-domain
307    sockets, this probably doesn't work on Windows.  It does work on most
308    posix platforms though (Linux and MacOS X, for example).
309    """
310
311    def __init__(self):
312        """
313        Open a session with the local machine's SSH agent, if one is running.
314        If no agent is running, initialization will succeed, but L{get_keys}
315        will return an empty tuple.
316
317        @raise SSHException: if an SSH agent is found, but speaks an
318            incompatible protocol
319        """
320        AgentSSH.__init__(self)
321
322        if ('SSH_AUTH_SOCK' in os.environ) and (sys.platform != 'win32'):
323            conn = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
324            try:
325                conn.connect(os.environ['SSH_AUTH_SOCK'])
326            except:
327                # probably a dangling env var: the ssh agent is gone
328                return
329        elif sys.platform == 'win32':
330            import win_pageant
331            if win_pageant.can_talk_to_agent():
332                conn = win_pageant.PageantConnection()
333            else:
334                return
335        else:
336            # no agent support
337            return
338        self._connect(conn)
339
340    def close(self):
341        """
342        Close the SSH agent connection.
343        """
344        self._close()
345
346class AgentKey(PKey):
347    """
348    Private key held in a local SSH agent.  This type of key can be used for
349    authenticating to a remote server (signing).  Most other key operations
350    work as expected.
351    """
352
353    def __init__(self, agent, blob):
354        self.agent = agent
355        self.blob = blob
356        self.name = Message(blob).get_string()
357
358    def __str__(self):
359        return self.blob
360
361    def get_name(self):
362        return self.name
363
364    def sign_ssh_data(self, rng, data):
365        msg = Message()
366        msg.add_byte(chr(SSH2_AGENTC_SIGN_REQUEST))
367        msg.add_string(self.blob)
368        msg.add_string(data)
369        msg.add_int(0)
370        ptype, result = self.agent._send_message(msg)
371        if ptype != SSH2_AGENT_SIGN_RESPONSE:
372            raise SSHException('key cannot be used for signing')
373        return result.get_string()
374