diff options
| -rw-r--r-- | paramiko/channel.py | 41 | ||||
| -rw-r--r-- | paramiko/client.py | 9 | ||||
| -rw-r--r-- | tests/test_client.py | 45 |
3 files changed, 93 insertions, 2 deletions
diff --git a/paramiko/channel.py b/paramiko/channel.py index 3a05bdc4..7735e1f1 100644 --- a/paramiko/channel.py +++ b/paramiko/channel.py @@ -283,6 +283,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 ebf21b08..681760cf 100644 --- a/paramiko/client.py +++ b/paramiko/client.py @@ -398,7 +398,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 @@ -411,6 +412,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 @@ -421,6 +423,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) @@ -428,7 +431,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 @@ -440,12 +443,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 d39febac..e7ebbc6a 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): @@ -373,3 +383,38 @@ class SSHClientTest (unittest.TestCase): password='pygmalion', ) self._test_connection(**kwargs) + + 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.') |