ssh: Use paramiko instead of pexpect

Author: DouglasRaillard

devlib/utils/ssh.py

@@ -29,6 +29,11 @@
 from pipes import quote
 from future.utils import raise_from
 
+from paramiko.client import SSHClient, AutoAddPolicy
+from paramiko.ssh_exception import NoValidConnectionsError
+# By default paramiko is very verbose, including at the INFO level
+logging.getLogger("paramiko").setLevel(logging.WARNING)
+
 # pylint: disable=import-error,wrong-import-position,ungrouped-imports,wrong-import-order
 import pexpect
 from distutils.version import StrictVersion as V
@@ -54,15 +59,7 @@
 logger = logging.getLogger('ssh')
 gem5_logger = logging.getLogger('gem5-connection')
 
-def ssh_get_shell(host,
-                  username,
-                  password=None,
-                  keyfile=None,
-                  port=None,
-                  timeout=10,
-                  telnet=False,
-                  original_prompt=None,
-                  options=None):
+def ssh_get_shell(host, username, password=None, keyfile=None, port=None, timeout=10, telnet=False, original_prompt=None):
     _check_env()
     start_time = time.time()
     while True:
@@ -71,8 +68,7 @@ def ssh_get_shell(host,
                 raise ValueError('keyfile may not be used with a telnet connection.')
             conn = TelnetPxssh(original_prompt=original_prompt)
         else:  # ssh
-            conn = pxssh.pxssh(options=options,
-                               echo=False)
+            conn = pxssh.pxssh(echo=False)
 
         try:
             if keyfile:
@@ -159,7 +155,6 @@ def check_keyfile(keyfile):
 
 class SshConnection(object):
 
-    default_password_prompt = '[sudo] password'
     max_cancel_attempts = 5
     default_timeout = 10
 
@@ -171,7 +166,7 @@ def name(self):
     def connected_as_root(self):
         if self._connected_as_root is None:
             # Execute directly to prevent deadlocking of connection
-            result = self._execute_and_wait_for_prompt('id', as_root=False)
+            exit_code, result = self._execute('id', as_root=False)
             self._connected_as_root = 'uid=0(' in result
         return self._connected_as_root
 
@@ -187,12 +182,8 @@ def __init__(self,
                  keyfile=None,
                  port=None,
                  timeout=None,
-                 telnet=False,
-                 password_prompt=None,
-                 original_prompt=None,
                  platform=None,
-                 sudo_cmd="sudo -- sh -c {}",
-                 options=None
+                 sudo_cmd="sudo -S -- sh -c {}"
                  ):
         self._connected_as_root = None
         self.host = host
@@ -201,185 +192,120 @@ def __init__(self,
         self.keyfile = check_keyfile(keyfile) if keyfile else keyfile
         self.port = port
         self.lock = threading.Lock()
-        self.password_prompt = password_prompt if password_prompt is not None else self.default_password_prompt
         self.sudo_cmd = sanitize_cmd_template(sudo_cmd)
         logger.debug('Logging in {}@{}'.format(username, host))
         timeout = timeout if timeout is not None else self.default_timeout
-        self.options = options if options is not None else {}
-        self.conn = ssh_get_shell(host,
-                                  username,
-                                  password,
-                                  self.keyfile,
-                                  port,
-                                  timeout,
-                                  False,
-                                  None,
-                                  self.options)
+
+        client = SSHClient()
+        client.load_system_host_keys()
+        client.set_missing_host_key_policy(AutoAddPolicy)
+        client.connect(
+            hostname=host,
+            port=port,
+            username=username,
+            password=password,
+            key_filename=keyfile,
+            timeout=timeout,
+        )
+        self.client = client
+        self.chan = self._get_channel()
         atexit.register(self.close)
 
+    def _get_channel(self):
+        transport = self.client.get_transport()
+        channel = transport.open_session()
+        return channel
+
+    def _get_sftp(self, timeout):
+        sftp = self.client.open_sftp()
+        sftp.get_channel().settimeout(timeout)
+        return sftp
+
     def push(self, source, dest, timeout=30):
-        dest = '{}@{}:{}'.format(self.username, self.host, dest)
-        return self._scp(source, dest, timeout)
+        sftp = self._get_sftp(timeout)
+        sftp.put(source, dest)
 
     def pull(self, source, dest, timeout=30):
-        source = '{}@{}:{}'.format(self.username, self.host, source)
-        return self._scp(source, dest, timeout)
+        sftp = self._get_sftp(timeout)
+        sftp.get(source, dest)
 
     def execute(self, command, timeout=None, check_exit_code=True,
                 as_root=False, strip_colors=True, will_succeed=False): #pylint: disable=unused-argument
         if command == '':
-            # Empty command is valid but the __devlib_ec stuff below will
-            # produce a syntax error with bash. Treat as a special case.
             return ''
         try:
+            _command = '({}) 2>&1'.format(command)
             with self.lock:
-                _command = '({}); __devlib_ec=$?; echo; echo $__devlib_ec'.format(command)
-                full_output = self._execute_and_wait_for_prompt(_command, timeout, as_root, strip_colors)
-                split_output = full_output.rsplit('\r\n', 2)
-                try:
-                    output, exit_code_text, _ = split_output
-                except ValueError as e:
-                    raise TargetStableError(
-                        "cannot split reply (target misconfiguration?):\n'{}'".format(full_output))
-                if check_exit_code:
-                    try:
-                        exit_code = int(exit_code_text)
-                        if exit_code:
-                            message = 'Got exit code {}\nfrom: {}\nOUTPUT: {}'
-                            raise TargetStableError(message.format(exit_code, command, output))
-                    except (ValueError, IndexError):
-                        logger.warning(
-                            'Could not get exit code for "{}",\ngot: "{}"'\
-                            .format(command, exit_code_text))
-                return output
-        except EOF:
+                exit_code, output = self._execute(_command, timeout, as_root, strip_colors)
+        except NoValidConnectionsError:
             raise TargetNotRespondingError('Connection lost.')
         except TargetStableError as e:
             if will_succeed:
                 raise TargetTransientError(e)
             else:
                 raise
+        else:
+            if check_exit_code and exit_code:
+                message = 'Got exit code {}\nfrom: {}\nOUTPUT: {}'
+                raise TargetStableError(message.format(exit_code, command, output))
+            return output
 
     def background(self, command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, as_root=False):
+        channel = self._get_channel()
         try:
-            port_string = '-p {}'.format(self.port) if self.port else ''
-            keyfile_string = '-i {}'.format(self.keyfile) if self.keyfile else ''
-            if as_root and not self.connected_as_root:
-                command = self.sudo_cmd.format(command)
-            options = " ".join([ "-o {}={}".format(key,val)
-                                for key,val in self.options.items()])
-            command = '{} {} {} {} {}@{} {}'.format(ssh,
-                                                    options,
-                                                    keyfile_string,
-                                                    port_string,
-                                                    self.username,
-                                                    self.host,
-                                                    command)
-            logger.debug(command)
-            if self.password:
-                command, _ = _give_password(self.password, command)
-            return subprocess.Popen(command, stdout=stdout, stderr=stderr, shell=True)
-        except EOF:
+            return channel.exec_command(command)
+        except NoValidConnectionsError:
             raise TargetNotRespondingError('Connection lost.')
 
     def close(self):
         logger.debug('Logging out {}@{}'.format(self.username, self.host))
         try:
-            self.conn.logout()
-        except:
+            self.client.close()
+        except NoValidConnectionsError:
             logger.debug('Connection lost.')
-            self.conn.close(force=True)
-
-    def cancel_running_command(self):
-        # simulate impatiently hitting ^C until command prompt appears
-        logger.debug('Sending ^C')
-        for _ in range(self.max_cancel_attempts):
-            self._sendline(chr(3))
-            if self.conn.prompt(0.1):
-                return True
-        return False
-
-    def _execute_and_wait_for_prompt(self, command, timeout=None, as_root=False, strip_colors=True, log=True):
-        self.conn.prompt(0.1)  # clear an existing prompt if there is one.
-        if as_root and self.connected_as_root:
-            # As we're already root, there is no need to use sudo.
-            as_root = False
-        if as_root:
-            command = self.sudo_cmd.format(quote(command))
-            if log:
-                logger.debug(command)
-            self._sendline(command)
-            if self.password:
-                index = self.conn.expect_exact([self.password_prompt, TIMEOUT], timeout=0.5)
-                if index == 0:
-                    self._sendline(self.password)
-        else:  # not as_root
-            if log:
-                logger.debug(command)
-            self._sendline(command)
-        timed_out = self._wait_for_prompt(timeout)
-        if sys.version_info[0] == 3:
-            output = process_backspaces(self.conn.before.decode(sys.stdout.encoding or 'utf-8', 'replace'))
-        else:
-            output = process_backspaces(self.conn.before)
 
-        if timed_out:
-            self.cancel_running_command()
-            raise TimeoutError(command, output)
+    def _execute(self, command, timeout=None, as_root=False, strip_colors=True, log=True):
+        # As we're already root, there is no need to use sudo.
+        use_sudo = not (as_root and self.connected_as_root)
+        log_debug = logger.debug if log else lambda msg: None
+
+        if use_sudo and not self.password:
+            raise TargetStableError('Attempt to use sudo but no password was specified')
+
+        try:
+            if use_sudo:
+                command = self.sudo_cmd.format(quote(command))
+                log_debug(command)
+                stdin, stdout, stderr = self.client.exec_command(command, timeout=timeout)
+                stdin.write(self.password + '\n')
+                stdin.flush()
+            else:
+                log_debug(command)
+                stdin, stdout, stderr = self.client.exec_command(command, timeout=timeout)
+        except socket.timeout:
+            raise TimeoutError(command)
+
+        # Empty the stdout buffer of the command, allowing it to carry on to
+        # completion
+        output_chunks = []
+        chunk = True
+        while chunk:
+            chunk = stdout.read()
+            output_chunks.append(chunk)
+
+        # Wait until the command completes
+        exit_code = stdout.channel.recv_exit_status()
+
+        # Join in one go to avoid O(N^2) concatenation
+        output = b''.join(output_chunks)
+
+        if sys.version_info[0] == 3:
+            output = output.decode(sys.stdout.encoding or 'utf-8', 'replace')
         if strip_colors:
             output = strip_bash_colors(output)
-        return output
 
-    def _wait_for_prompt(self, timeout=None):
-        if timeout:
-            return not self.conn.prompt(timeout)
-        else:  # cannot timeout; wait forever
-            while not self.conn.prompt(1):
-                pass
-            return False
-
-    def _scp(self, source, dest, timeout=30):
-        # NOTE: the version of scp in Ubuntu 12.04 occasionally (and bizarrely)
-        # fails to connect to a device if port is explicitly specified using -P
-        # option, even if it is the default port, 22. To minimize this problem,
-        # only specify -P for scp if the port is *not* the default.
-        port_string = '-P {}'.format(quote(str(self.port))) if (self.port and self.port != 22) else ''
-        keyfile_string = '-i {}'.format(quote(self.keyfile)) if self.keyfile else ''
-        options = " ".join(["-o {}={}".format(key,val)
-                            for key,val in self.options.items()])
-        command = '{} {} -r {} {} {} {}'.format(scp,
-                                                options,
-                                                keyfile_string,
-                                                port_string,
-                                                quote(source),
-                                                quote(dest))
-        command_redacted = command
-        logger.debug(command)
-        if self.password:
-            command, command_redacted = _give_password(self.password, command)
-        try:
-            check_output(command, timeout=timeout, shell=True)
-        except subprocess.CalledProcessError as e:
-            raise_from(HostError("Failed to copy file with '{}'. Output:\n{}".format(
-                command_redacted, e.output)), None)
-        except TimeoutError as e:
-            raise TimeoutError(command_redacted, e.output)
-
-    def _sendline(self, command):
-        # Workaround for https://github.com/pexpect/pexpect/issues/552
-        if len(command) == self._get_window_size()[1] - self._get_prompt_length():
-            command += ' '
-        self.conn.sendline(command)
-
-    @memoized
-    def _get_prompt_length(self):
-        self.conn.sendline()
-        self.conn.prompt()
-        return len(self.conn.after)
+        return (exit_code, output)
 
-    @memoized
-    def _get_window_size(self):
-        return self.conn.getwinsize()
 
 class TelnetConnection(SshConnection):
 

setup.py

@@ -82,6 +82,7 @@
         'python-dateutil',  # converting between UTC and local time.
         'pexpect>=3.3',  # Send/recieve to/from device
         'pyserial',  # Serial port interface
+        'paramiko', # SSH connection
         'wrapt',  # Basic for construction of decorator functions
         'future', # Python 2-3 compatibility
         'enum34;python_version<"3.4"', # Enums for Python < 3.4