Support transmission of environment variables

paramiko/channel.py

@@ -278,6 +278,47 @@ def resize_pty(self, width=80, height=24, width_pixels=0, height_pixels=0):
         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

paramiko/client.py

@@ -319,7 +319,8 @@ def close(self):
             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 @@ def exec_command(self, command, bufsize=-1, timeout=None, get_pty=False):
             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,14 +344,15 @@ def exec_command(self, command, bufsize=-1, timeout=None, get_pty=False):
         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)
         stderr = chan.makefile_stderr('r', bufsize)
         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 @@ def invoke_shell(self, term='vt100', width=80, height=24, width_pixels=0,
         :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
 

tests/test_client.py

@@ -79,6 +79,16 @@ def check_channel_exec_request(self, channel, command):
             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 @@ def test_7_banner_timeout(self):
             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.')