Support transmission of environment variables

The SSH protocol allows the client to transmit environment variables to
the server. This is particularly useful if the user wants to modify the
environment of an executed command without having to reexecute the
actual command from a shell.

This patch extends the Client and Channel interface to allow
the transmission of environment variables to the server side.

In order to use this feature the SSH server must accept environment
variables from the client (e.g. the AcceptEnv configuration directive of
OpenSSH).

3 files changed, 93 insertions, 2 deletions

diff --git a/paramiko/channel.py b/paramiko/channel.py
index 9de278cb..01b815e5 100644
--- a/paramiko/channel.py
+++ b/paramiko/channel.py
@@ -278,6 +278,47 @@ class Channel (ClosingContextManager):
         m.add_int(height_pixels)
         self.transport._send_user_message(m)
 
+    @open_only
+    def update_environment_variables(self, environment):
+        """
+        Updates this channel's environment. This operation is additive - i.e.
+        the current environment is not reset before the given environment
+        variables are set.
+
+        :param dict environment: a dictionary containing the name and respective
+            values to set
+        :raises SSHException:
+            if any of the environment variables was rejected by the server or
+            the channel was closed
+        """
+        for name, value in environment.items():
+            try:
+                self.set_environment_variable(name, value)
+            except SSHException as e:
+                raise SSHException("Failed to set environment variable \"%s\"." % name, e)
+
+    @open_only
+    def set_environment_variable(self, name, value):
+        """
+        Set the value of an environment variable.
+
+        :param str name: name of the environment variable
+        :param str value: value of the environment variable
+
+        :raises SSHException:
+            if the request was rejected or the channel was closed
+        """
+        m = Message()
+        m.add_byte(cMSG_CHANNEL_REQUEST)
+        m.add_int(self.remote_chanid)
+        m.add_string('env')
+        m.add_boolean(True)
+        m.add_string(name)
+        m.add_string(value)
+        self._event_pending()
+        self.transport._send_user_message(m)
+        self._wait_for_event()
+
     def exit_status_ready(self):
         """
         Return true if the remote process has exited and returned an exit
diff --git a/paramiko/client.py b/paramiko/client.py
index 05686d97..99bce156 100644
--- a/paramiko/client.py
+++ b/paramiko/client.py
@@ -319,7 +319,8 @@ class SSHClient (ClosingContextManager):
             self._agent.close()
             self._agent = None
 
-    def exec_command(self, command, bufsize=-1, timeout=None, get_pty=False):
+    def exec_command(self, command, bufsize=-1, timeout=None, get_pty=False,
+                     environment=None):
         """
         Execute a command on the SSH server.  A new `.Channel` is opened and
         the requested command is executed.  The command's input and output
@@ -332,6 +333,7 @@ class SSHClient (ClosingContextManager):
             Python
         :param int timeout:
             set command's channel timeout. See `Channel.settimeout`.settimeout
+        :param dict environment: the command's environment
         :return:
             the stdin, stdout, and stderr of the executing command, as a
             3-tuple
@@ -342,6 +344,7 @@ class SSHClient (ClosingContextManager):
         if get_pty:
             chan.get_pty()
         chan.settimeout(timeout)
+        chan.update_environment_variables(environment or {})
         chan.exec_command(command)
         stdin = chan.makefile('wb', bufsize)
         stdout = chan.makefile('r', bufsize)
@@ -349,7 +352,7 @@ class SSHClient (ClosingContextManager):
         return stdin, stdout, stderr
 
     def invoke_shell(self, term='vt100', width=80, height=24, width_pixels=0,
-                     height_pixels=0):
+                     height_pixels=0, environment=None):
         """
         Start an interactive shell session on the SSH server.  A new `.Channel`
         is opened and connected to a pseudo-terminal using the requested
@@ -361,12 +364,14 @@ class SSHClient (ClosingContextManager):
         :param int height: the height (in characters) of the terminal window
         :param int width_pixels: the width (in pixels) of the terminal window
         :param int height_pixels: the height (in pixels) of the terminal window
+        :param dict environment: the command's environment
         :return: a new `.Channel` connected to the remote shell
 
         :raises SSHException: if the server fails to invoke a shell
         """
         chan = self._transport.open_session()
         chan.get_pty(term, width, height, width_pixels, height_pixels)
+        chan.update_environment_variables(environment or {})
         chan.invoke_shell()
         return chan
 
diff --git a/tests/test_client.py b/tests/test_client.py
index 28d1cb46..0022a4d9 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -79,6 +79,16 @@ class NullServer (paramiko.ServerInterface):
             return False
         return True
 
+    def check_channel_env_request(self, channel, name, value):
+        if name == 'INVALID_ENV':
+            return False
+
+        if not hasattr(channel, 'env'):
+            setattr(channel, 'env', {})
+
+        channel.env[name] = value
+        return True
+
 
 class SSHClientTest (unittest.TestCase):
 
@@ -344,3 +354,38 @@ class SSHClientTest (unittest.TestCase):
             password='pygmalion',
             banner_timeout=0.5
         )
+
+    def test_update_environment(self):
+        """
+        Verify that environment variables can be set by the client.
+        """
+        threading.Thread(target=self._run).start()
+
+        self.tc = paramiko.SSHClient()
+        self.tc.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+        self.assertEqual(0, len(self.tc.get_host_keys()))
+        self.tc.connect(self.addr, self.port, username='slowdive', password='pygmalion')
+
+        self.event.wait(1.0)
+        self.assertTrue(self.event.isSet())
+        self.assertTrue(self.ts.is_active())
+
+        target_env = {b'A': b'B', b'C': b'd'}
+
+        self.tc.exec_command('yes', environment=target_env)
+        schan = self.ts.accept(1.0)
+        self.assertEqual(target_env, getattr(schan, 'env', {}))
+        schan.close()
+
+        # Cannot use assertRaises in context manager mode as it is not supported
+        # in Python 2.6.
+        try:
+            # Verify that a rejection by the server can be detected
+            self.tc.exec_command('yes', environment={b'INVALID_ENV': b''})
+        except SSHException as e:
+            self.assertTrue('INVALID_ENV' in str(e),
+                            'Expected variable name in error message')
+            self.assertTrue(isinstance(e.args[1], SSHException),
+                            'Expected original SSHException in exception')
+        else:
+            self.assertFalse(False, 'SSHException was not thrown.')