Made output decoding codec configurable at run_command, added test

Author: Dan

embedded_server/embedded_server.py

@@ -103,7 +103,8 @@ class Server(paramiko.ServerInterface):
     """
 
     def __init__(self, transport, host_key, fail_auth=False,
-                 ssh_exception=False):
+                 ssh_exception=False,
+                 encoding='utf-8'):
         paramiko.ServerInterface.__init__(self)
         transport.load_server_moduli()
         transport.add_server_key(host_key)
@@ -113,6 +114,7 @@ def __init__(self, transport, host_key, fail_auth=False,
         self.fail_auth = fail_auth
         self.ssh_exception = ssh_exception
         self.host_key = host_key
+        self.encoding = encoding
 
     def check_channel_request(self, kind, chanid):
         return paramiko.OPEN_SUCCEEDED
@@ -163,12 +165,11 @@ def check_channel_forward_agent_request(self, channel):
         gevent.sleep()
         return True
 
-    def check_channel_exec_request(self, channel, cmd,
-                                   encoding='utf-8'):
+    def check_channel_exec_request(self, channel, cmd):
         logger.debug("Got exec request on channel %s for cmd %s" % (channel, cmd,))
         self.event.set()
         _env = os.environ
-        _env['PYTHONIOENCODING'] = encoding
+        _env['PYTHONIOENCODING'] = self.encoding
         if hasattr(channel, 'environment'):
             _env.update(channel.environment)
         process = gevent.subprocess.Popen(cmd, stdout=gevent.subprocess.PIPE,
@@ -182,7 +183,8 @@ def check_channel_exec_request(self, channel, cmd,
     def check_channel_env_request(self, channel, name, value):
         if not hasattr(channel, 'environment'):
             channel.environment = {}
-        channel.environment.update({name.decode('utf8'): value.decode('utf8')})
+        channel.environment.update({
+            name.decode(self.encoding): value.decode(self.encoding)})
         return True
 
     def _read_response(self, channel, process):
@@ -213,7 +215,8 @@ def make_socket(listen_ip, port=0):
     return sock
 
 def listen(sock, fail_auth=False, ssh_exception=False,
-           timeout=None):
+           timeout=None,
+           encoding='utf-8'):
     """Run server and given a cmd_to_run, send given
     response to client connection. Returns (server, socket) tuple
     where server is a joinable server thread and socket is listening
@@ -228,13 +231,16 @@ def listen(sock, fail_auth=False, ssh_exception=False,
         return
     host, port = sock.getsockname()
     logger.info('Listening for connection on %s:%s..', host, port)
-    return handle_ssh_connection(sock, fail_auth=fail_auth,
-                          timeout=timeout, ssh_exception=ssh_exception)
+    return handle_ssh_connection(
+        sock, fail_auth=fail_auth, timeout=timeout, ssh_exception=ssh_exception,
+        encoding=encoding)
 
 def _handle_ssh_connection(transport, fail_auth=False,
-                           ssh_exception=False):
+                           ssh_exception=False,
+                           encoding='utf-8'):
     server = Server(transport, host_key,
-                    fail_auth=fail_auth, ssh_exception=ssh_exception)
+                    fail_auth=fail_auth, ssh_exception=ssh_exception,
+                    encoding=encoding)
     try:
         transport.start_server(server=server)
     except paramiko.SSHException as e:
@@ -261,7 +267,8 @@ def _handle_ssh_connection(transport, fail_auth=False,
 
 def handle_ssh_connection(sock,
                           fail_auth=False, ssh_exception=False,
-                          timeout=None):
+                          timeout=None,
+                          encoding='utf-8'):
     conn, addr = sock.accept()
     logger.info('Got connection..')
     if timeout:
@@ -271,7 +278,8 @@ def handle_ssh_connection(sock,
     try:
         transport = paramiko.Transport(conn)
         return _handle_ssh_connection(transport, fail_auth=fail_auth,
-                                      ssh_exception=ssh_exception)
+                                      ssh_exception=ssh_exception,
+                                      encoding=encoding)
     except Exception as e:
         logger.error('*** Caught exception: %s: %s' % (str(e.__class__), str(e),))
         traceback.print_exc()
@@ -281,14 +289,18 @@ def handle_ssh_connection(sock,
             pass
 
 def start_server(sock, fail_auth=False, ssh_exception=False,
-                 timeout=None):
+                 timeout=None,
+                 encoding='utf-8'):
     return gevent.spawn(listen, sock, fail_auth=fail_auth,
-                        timeout=timeout, ssh_exception=ssh_exception)
+                        timeout=timeout, ssh_exception=ssh_exception,
+                        encoding=encoding)
 
 def start_server_from_ip(ip, port=0,
                          fail_auth=False, ssh_exception=False,
-                         timeout=None):
+                         timeout=None,
+                         encoding='utf-8'):
     server_sock = make_socket(ip, port=port)
     server = start_server(server_sock, fail_auth=fail_auth,
-                          ssh_exception=ssh_exception, timeout=timeout)
+                          ssh_exception=ssh_exception, timeout=timeout,
+                          encoding=encoding)
     return server, server_sock.getsockname()[1]

pssh/pssh_client.py

@@ -336,7 +336,8 @@ def __init__(self, hosts,
         self.channel_timeout = channel_timeout
 
     def run_command(self, command, sudo=False, user=None, stop_on_errors=True,
-                    shell=None, use_shell=True, use_pty=True, host_args=None):
+                    shell=None, use_shell=True, use_pty=True, host_args=None,
+                    encoding='utf-8'):
         """Run command on all hosts in parallel, honoring self.pool_size,
         and return output buffers.
 
@@ -383,6 +384,9 @@ def run_command(self, command, sudo=False, user=None, stop_on_errors=True,
         host list - :py:class:`pssh.exceptions.HostArgumentException` is raised \
         otherwise
         :type host_args: tuple or list
+        :param encoding: Encoding to use for output. Must be valid
+            `Python codec <https://docs.python.org/2.7/library/codecs.html>`_
+        :type encoding: str
 
         :rtype: Dictionary with host as key and \
         :py:class:`pssh.output.HostOutput` as value as per \
@@ -639,7 +643,7 @@ def run_command(self, command, sudo=False, user=None, stop_on_errors=True,
                 for host in self.hosts]
         for cmd in cmds:
             try:
-                self.get_output(cmd, output)
+                self.get_output(cmd, output, encoding=encoding)
             except Exception:
                 if stop_on_errors:
                     raise
@@ -676,7 +680,7 @@ def _exec_command(self, host, command, sudo=False, user=None,
             command, sudo=sudo, user=user, shell=shell,
             use_shell=use_shell, use_pty=use_pty)
 
-    def get_output(self, cmd, output):
+    def get_output(self, cmd, output, encoding='utf-8'):
         """Get output from command.
 
         :param cmd: Command to get output from
@@ -733,10 +737,12 @@ def get_output(self, cmd, output):
             raise
         stdout = self.host_clients[host].read_output_buffer(
             stdout, callback=self.get_exit_codes,
-            callback_args=(output,))
+            callback_args=(output,),
+            encoding=encoding)
         stderr = self.host_clients[host].read_output_buffer(
             stderr, prefix='\t[err]', callback=self.get_exit_codes,
-            callback_args=(output,))
+            callback_args=(output,),
+            encoding=encoding)
         self._update_host_output(output, host, self._get_exit_code(channel),
                                  channel, stdout, stderr, stdin, cmd)
 

pssh/ssh_client.py

@@ -262,7 +262,8 @@ def exec_command(self, command, sudo=False, user=None,
 
     def read_output_buffer(self, output_buffer, prefix='',
                            callback=None,
-                           callback_args=None):
+                           callback_args=None,
+                           encoding='utf-8'):
         """Read from output buffers and log to host_logger
 
         :param output_buffer: Iterator containing buffer
@@ -274,7 +275,7 @@ def read_output_buffer(self, output_buffer, prefix='',
         :param callback_args: Arguments for call back function
         :type callback_args: tuple"""
         for line in output_buffer:
-            output = line.strip().decode('utf8')
+            output = line.strip().decode(encoding)
             host_logger.info("[%s]%s\t%s", self.host, prefix, output,)
             yield output
         if callback:

tests/test_pssh_client.py

@@ -947,6 +947,16 @@ def test_ssh_client_utf_encoding(self):
         self.assertEqual(expected, stdout,
                          msg="Got unexpected unicode output %s - expected %s" % (
                              stdout, expected,))
+        utf16_server, server_port = start_server_from_ip(
+            self.host, encoding='utf-16')
+        client = ParallelSSHClient([self.host], port=server_port,
+                                   pkey=self.user_key)
+        # File is already set to utf-8, cannot use utf-16 only representations
+        # Using ascii characters encoded as utf-16 instead
+        output = client.run_command(self.fake_cmd, encoding='utf-16')
+        stdout = list(output[self.host]['stdout'])
+        # import ipdb; ipdb.set_trace()
+        self.assertEqual([self.fake_resp.decode('utf-16')], stdout)
 
     def test_pty(self):
         cmd = "exit 0"