Added ssh agent support. Ported from https://github.com/robey/paramiko/pull/21

(cherry picked from commit 35a173631f4d5be2e47d8880ee0a1df08f95cebc)

Conflicts:
	paramiko/agent.py

4 files changed, 333 insertions, 52 deletions

diff --git a/paramiko/agent.py b/paramiko/agent.py
index 3bb94261..4b562f78 100644
--- a/paramiko/agent.py
+++ b/paramiko/agent.py
@@ -24,17 +24,265 @@ import os
 import socket
 import struct
 import sys
+import threading
+import tempfile
+import stat
+import select
+import fcntl
 
 from paramiko.ssh_exception import SSHException
 from paramiko.message import Message
 from paramiko.pkey import PKey
-
+from paramiko.channel import Channel
 
 SSH2_AGENTC_REQUEST_IDENTITIES, SSH2_AGENT_IDENTITIES_ANSWER, \
     SSH2_AGENTC_SIGN_REQUEST, SSH2_AGENT_SIGN_RESPONSE = range(11, 15)
 
+class AgentSSH:
+    """
+    Client interface for using private keys from an SSH agent running on the
+    local machine.  If an SSH agent is running, this class can be used to
+    connect to it and retreive L{PKey} objects which can be used when
+    attempting to authenticate to remote SSH servers.
+    
+    Because the SSH agent protocol uses environment variables and unix-domain
+    sockets, this probably doesn't work on Windows.  It does work on most
+    posix platforms though (Linux and MacOS X, for example).
+    """
+    def __init__(self):
+        self._conn = None
+        self._keys = ()
+
+    def get_keys(self):
+        """
+        Return the list of keys available through the SSH agent, if any.  If
+        no SSH agent was running (or it couldn't be contacted), an empty list
+        will be returned.
+        
+        @return: a list of keys available on the SSH agent
+        @rtype: tuple of L{AgentKey}
+        """
+        return self._keys
+
+    def _connect(self, conn):
+        self._conn = conn
+        ptype, result = self._send_message(chr(SSH2_AGENTC_REQUEST_IDENTITIES))
+        if ptype != SSH2_AGENT_IDENTITIES_ANSWER:
+            raise SSHException('could not get keys from ssh-agent')
+        keys = []
+        for i in range(result.get_int()):
+            keys.append(AgentKey(self, result.get_string()))
+            result.get_string()
+        self._keys = tuple(keys)
+
+    def _close(self):
+        #self._conn.close()
+        self._conn = None
+        self._keys = ()
+
+    def _send_message(self, msg):
+        msg = str(msg)
+        self._conn.send(struct.pack('>I', len(msg)) + msg)
+        l = self._read_all(4)
+        msg = Message(self._read_all(struct.unpack('>I', l)[0]))
+        return ord(msg.get_byte()), msg
+
+    def _read_all(self, wanted):
+        result = self._conn.recv(wanted)
+        while len(result) < wanted:
+            if len(result) == 0:
+                raise SSHException('lost ssh-agent')
+            extra = self._conn.recv(wanted - len(result))
+            if len(extra) == 0:
+                raise SSHException('lost ssh-agent')
+            result += extra
+        return result
+
+class AgentProxyThread(threading.Thread):
+    """ Class in charge of communication between two chan """
+    def __init__(self, agent):
+        threading.Thread.__init__(self, target=self.run)
+        self._agent = agent
+        self._exit = False
+
+    def run(self):
+        try:
+            (r,addr) = self.get_connection()
+            self.__inr = r
+            self.__addr = addr
+            self._agent.connect()
+            self._communicate()
+        except:
+            #XXX Not sure what to do here ... raise or pass ?
+            raise
+
+    def _communicate(self):
+        p = select.poll()
+        oldflags = fcntl.fcntl(self.__inr, fcntl.F_GETFL)
+        fcntl.fcntl(self.__inr, fcntl.F_SETFL, oldflags | os.O_NONBLOCK)
+        p.register(self._agent._conn, select.POLLIN)
+        p.register(self.__inr, select.POLLIN)
+        while not self._exit:
+            c = p.poll(500)
+            for cc in c:
+                fd, event = cc
+                if self._agent._conn.fileno() == fd:
+                    data = self._agent._conn.recv(512)
+                    if len(data) != 0:
+                        self.__inr.send(data)
+                    else:
+                        break
+                elif self.__inr.fileno() == fd:
+                    data = self.__inr.recv(512)
+                    if len(data) != 0:
+                        self._agent._conn.send(data)
+                    else:
+                        break
+
+class AgentLocalProxy(AgentProxyThread):
+    """
+    Class to be used when wanting to ask a local SSH Agent being
+    asked from a remote fake agent (so use a unix socket for ex.)
+    """
+    def __init__(self, agent):
+        AgentProxyThread.__init__(self, agent)
+
+    def get_connection(self):
+        """ Return a pair of socket object and string address
+        May Block !
+        """
+        conn = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+        try:
+            conn.bind(self._agent._get_filename())
+            conn.listen(1)
+            (r,addr) = conn.accept()
+            return (r, addr)
+        except:
+            raise
+        return None
+
+class AgentRemoteProxy(AgentProxyThread):
+    """
+    Class to be used when wanting to ask a remote SSH Agent
+    """
+    def __init__(self, agent, chan):
+        AgentProxyThread.__init__(self, agent)
+        self.__chan = chan
+
+    def get_connection(self):
+        """
+        Class to be used when wanting to ask a local SSH Agent being
+        asked from a remote fake agent (so use a unix socket for ex.)
+        """
+        return (self.__chan, None)
 
-class Agent:
+class AgentClientProxy:
+    """
+    Class proxying request as a client:
+       -> client ask for a request_forward_agent()
+       -> server creates a proxy and a fake SSH Agent
+       -> server ask for establishing a connection when needed,
+       calling the forward_agent_handler at client side.
+       -> the forward_agent_handler launch a thread for connecting
+       the remote fake agent and the local agent
+       -> Communication occurs ...
+    """
+    def __init__(self, chanClient):
+        self._conn = None
+        self.__chanC = chanClient
+        chanClient.request_forward_agent(self._forward_agent_handler)
+
+    def __del__(self):
+        self.close()
+
+    def connect(self):
+        """ 
+        Method automatically called by the run() method of the AgentProxyThread
+        """
+        if ('SSH_AUTH_SOCK' in os.environ) and (sys.platform != 'win32'):
+            conn = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+            try:
+                conn.connect(os.environ['SSH_AUTH_SOCK'])
+            except:
+                # probably a dangling env var: the ssh agent is gone
+                return
+        elif sys.platform == 'win32':
+            import win_pageant
+            if win_pageant.can_talk_to_agent():
+                conn = win_pageant.PageantConnection()
+            else:
+                return
+        else:
+            # no agent support
+            return
+        self._conn = conn
+
+    def close(self):
+        """
+        Close the current connection and terminate the agent
+        Should be called manually
+        """
+        if hasattr(self, "thread"):
+            self.thread._exit = True
+            self.thread.join(1000)
+        if self._conn is not None:
+            self._conn.close()
+
+    def _forward_agent_handler(self, chanRemote):
+        self.thread = AgentRemoteProxy(self, chanRemote)
+        self.thread.start()
+
+class AgentServerProxy(AgentSSH):
+    """
+    @param t : transport used for the Forward for SSH Agent communication
+
+    @raise SSHException: mostly if we lost the agent
+    """
+    def __init__(self, t):
+        AgentSSH.__init__(self)
+        self.__t = t
+        self._dir = tempfile.mkdtemp('sshproxy')
+        os.chmod(self._dir, stat.S_IRWXU)
+        self._file = self._dir + '/sshproxy.ssh'
+        self.thread = AgentLocalProxy(self)
+        self.thread.start()
+
+    def __del__(self):
+        self.close()
+
+    def connect(self):
+         conn_sock = self.__t.open_forward_agent_channel()
+         if conn_sock is None:
+             raise SSHException('lost ssh-agent')
+         conn_sock.set_name('auth-agent')
+         self._connect(conn_sock)
+
+    def close(self):
+        """
+        Terminate the agent, clean the files, close connections
+        Should be called manually
+        """
+        os.remove(self._file)
+        os.rmdir(self._dir)
+        self.thread._exit = True
+        self.thread.join(1000)
+        self._close()
+
+    def get_env(self):
+        """
+        Helper for the environnement under unix
+
+        @return: the SSH_AUTH_SOCK Environnement variables
+        @rtype: dict
+        """
+        env = {} 
+        env['SSH_AUTH_SOCK'] = self._get_filename()
+        return env
+
+    def _get_filename(self):
+        return self._file
+
+class Agent(AgentSSH):
     """
     Client interface for using private keys from an SSH agent running on the
     local machine.  If an SSH agent is running, this class can be used to
@@ -55,8 +303,8 @@ class Agent:
         @raise SSHException: if an SSH agent is found, but speaks an
             incompatible protocol
         """
-        self.conn = None
-        self.keys = ()
+        AgentSSH.__init__(self)
+
         if ('SSH_AUTH_SOCK' in os.environ) and (sys.platform != 'win32'):
             conn = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
             try:
@@ -64,64 +312,22 @@ class Agent:
             except:
                 # probably a dangling env var: the ssh agent is gone
                 return
-            self.conn = conn
         elif sys.platform == 'win32':
             import win_pageant
             if win_pageant.can_talk_to_agent():
-                self.conn = win_pageant.PageantConnection()
+                conn = win_pageant.PageantConnection()
             else:
                 return
         else:
             # no agent support
             return
-            
-        ptype, result = self._send_message(chr(SSH2_AGENTC_REQUEST_IDENTITIES))
-        if ptype != SSH2_AGENT_IDENTITIES_ANSWER:
-            raise SSHException('could not get keys from ssh-agent')
-        keys = []
-        for i in range(result.get_int()):
-            keys.append(AgentKey(self, result.get_string()))
-            result.get_string()
-        self.keys = tuple(keys)
+        self._connect(conn)            
 
     def close(self):
         """
         Close the SSH agent connection.
         """
-        if self.conn is not None:
-            self.conn.close()
-        self.conn = None
-        self.keys = ()
-
-    def get_keys(self):
-        """
-        Return the list of keys available through the SSH agent, if any.  If
-        no SSH agent was running (or it couldn't be contacted), an empty list
-        will be returned.
-        
-        @return: a list of keys available on the SSH agent
-        @rtype: tuple of L{AgentKey}
-        """
-        return self.keys
-
-    def _send_message(self, msg):
-        msg = str(msg)
-        self.conn.send(struct.pack('>I', len(msg)) + msg)
-        l = self._read_all(4)
-        msg = Message(self._read_all(struct.unpack('>I', l)[0]))
-        return ord(msg.get_byte()), msg
-
-    def _read_all(self, wanted):
-        result = self.conn.recv(wanted)
-        while len(result) < wanted:
-            if len(result) == 0:
-                raise SSHException('lost ssh-agent')
-            extra = self.conn.recv(wanted - len(result))
-            if len(extra) == 0:
-                raise SSHException('lost ssh-agent')
-            result += extra
-        return result
-
+        self._close()
 
 class AgentKey(PKey):
     """
diff --git a/paramiko/channel.py b/paramiko/channel.py
index 6d895fe4..534f8d7c 100644
--- a/paramiko/channel.py
+++ b/paramiko/channel.py
@@ -381,6 +381,31 @@ class Channel (object):
         self.transport._set_x11_handler(handler)
         return auth_cookie
 
+    def request_forward_agent(self, handler):
+        """
+        Request for a forward SSH Agent on this channel.
+        This is only valid for an ssh-agent from openssh !!!
+
+        @param handler: a required handler to use for incoming SSH Agent connections
+        @type handler: function
+
+        @return: if we are ok or not (at that time we always return ok)
+        @rtype: boolean
+
+        @raise: SSHException in case of channel problem.
+        """
+        if self.closed or self.eof_received or self.eof_sent or not self.active:
+            raise SSHException('Channel is not open')
+
+        m = Message()
+        m.add_byte(chr(MSG_CHANNEL_REQUEST))
+        m.add_int(self.remote_chanid)
+        m.add_string('auth-agent-req@openssh.com')
+        m.add_boolean(False)
+        self.transport._send_user_message(m)
+        self.transport._set_forward_agent_handler(handler)
+        return True
+
     def get_transport(self):
         """
         Return the L{Transport} associated with this channel.
@@ -1026,6 +1051,11 @@ class Channel (object):
             else:
                 ok = server.check_channel_x11_request(self, single_connection,
                                                       auth_proto, auth_cookie, screen_number)
+        elif key == 'auth-agent-req@openssh.com':
+            if server is None:
+                ok = False
+            else:
+                ok = server.check_channel_forward_agent_request(self)
         else:
             self._log(DEBUG, 'Unhandled channel request "%s"' % key)
             ok = False
diff --git a/paramiko/server.py b/paramiko/server.py
index 6424b63a..2a1172aa 100644
--- a/paramiko/server.py
+++ b/paramiko/server.py
@@ -93,6 +93,7 @@ class ServerInterface (object):
             - L{check_channel_subsystem_request}
             - L{check_channel_window_change_request}
             - L{check_channel_x11_request}
+            - L{check_channel_forward_agent_request}
 
         The C{chanid} parameter is a small number that uniquely identifies the
         channel within a L{Transport}.  A L{Channel} object is not created
@@ -492,7 +493,21 @@ class ServerInterface (object):
         @rtype: bool
         """
         return False
-    
+
+    def check_channel_forward_agent_request(self, channel):
+        """
+	Determine if the client will be provided with an forward agent session. If this
+	method returns C{True}, the server will allow SSH Agent forwarding.
+
+	The default implementation always returns C{False}.
+
+	@param channel: the L{Channel} the request arrived on
+	@type channel: L{Channel}
+	@return: C{True} if the AgentForward was loaded; C{False} if not
+	@rtype: bool
+	"""
+	return False
+
     def check_channel_direct_tcpip_request(self, chanid, origin, destination):
         """
         Determine if a local port forwarding channel will be granted, and
diff --git a/paramiko/transport.py b/paramiko/transport.py
index 7cbeeea5..42a371f3 100644
--- a/paramiko/transport.py
+++ b/paramiko/transport.py
@@ -341,6 +341,7 @@ class Transport (threading.Thread):
         self._channel_counter = 1
         self.window_size = 65536
         self.max_packet_size = 34816
+        self._forward_agent_handler = None
         self._x11_handler = None
         self._tcp_handler = None
 
@@ -672,6 +673,18 @@ class Transport (threading.Thread):
             prematurely
         """
         return self.open_channel('x11', src_addr=src_addr)
+ 
+    def open_forward_agent_channel(self):
+         """
+	 Request a new channel to the client, of type C{"auth-agent@openssh.com"}.
+	 This is just an alias for C{open_channel('auth-agent@openssh.com')}.
+	 @return: a new L{Channel}
+	 @rtype: L{Channel}
+
+	 @raise SSHException: if the request is rejected or the session ends
+	     prematurely
+	 """
+         return self.open_channel('auth-agent@openssh.com')
 
     def open_forwarded_tcpip_channel(self, (src_addr, src_port), (dest_addr, dest_port)):
         """
@@ -1481,6 +1494,14 @@ class Transport (threading.Thread):
         else:
             return self._cipher_info[name]['class'].new(key, self._cipher_info[name]['mode'], iv)
 
+    def _set_forward_agent_handler(self, handler):
+        if handler is None:
+            def default_handler(channel):
+                self._queue_incoming_channel(channel)
+            self._forward_agent_handler = default_handler
+        else:
+            self._forward_agent_handler = handler
+
     def _set_x11_handler(self, handler):
         # only called if a channel has turned on x11 forwarding
         if handler is None:
@@ -1996,7 +2017,14 @@ class Transport (threading.Thread):
         initial_window_size = m.get_int()
         max_packet_size = m.get_int()
         reject = False
-        if (kind == 'x11') and (self._x11_handler is not None):
+	if (kind == 'auth-agent@openssh.com') and (self._forward_agent_handler is not None):
+	    self._log(DEBUG, 'Incoming forward agent connection')
+            self.lock.acquire()
+            try:
+                my_chanid = self._next_channel()
+            finally:
+                self.lock.release()
+        elif (kind == 'x11') and (self._x11_handler is not None):
             origin_addr = m.get_string()
             origin_port = m.get_int()
             self._log(DEBUG, 'Incoming x11 connection from %s:%d' % (origin_addr, origin_port))
@@ -2068,7 +2096,9 @@ class Transport (threading.Thread):
         m.add_int(self.max_packet_size)
         self._send_message(m)
         self._log(INFO, 'Secsh channel %d (%s) opened.', my_chanid, kind)
-        if kind == 'x11':
+	if kind == 'auth-agent@openssh.com':
+            self._forward_agent_handler(chan)
+        elif kind == 'x11':
             self._x11_handler(chan, (origin_addr, origin_port))
         elif kind == 'forwarded-tcpip':
             chan.origin_addr = (origin_addr, origin_port)