Source code for wlauto.utils.ssh

#    Copyright 2014-2015 ARM Limited
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import stat
import logging
import subprocess
import re
import threading
import tempfile
import shutil
import time

from pexpect import EOF, TIMEOUT, spawn, pxssh

from wlauto.exceptions import HostError, DeviceError, TimeoutError, ConfigError
from wlauto.utils.misc import (which, strip_bash_colors, escape_single_quotes, check_output,

ssh = None
scp = None
sshpass = None

logger = logging.getLogger('ssh')

[docs]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() extra_login_args = {} while True: if telnet: if keyfile: raise ValueError('keyfile may not be used with a telnet connection.') conn = TelnetConnection() if original_prompt: extra_login_args['original_prompt'] = original_prompt if port is None: port = 23 else: # ssh conn = pxssh.pxssh() try: if keyfile: conn.login(host, username, ssh_key=keyfile, port=port, login_timeout=timeout, **extra_login_args) else: conn.login(host, username, password, port=port, login_timeout=timeout, **extra_login_args) break except EOF: timeout -= time.time() - start_time if timeout <= 0: message = 'Could not connect to {}; is the host name correct?' raise DeviceError(message.format(host)) time.sleep(5) conn.sendline('shopt -s checkwinsize') conn.prompt() conn.setwinsize(500,200) conn.sendline('') conn.prompt() conn.sendline('stty rows 500') conn.prompt() conn.sendline('stty cols 200') conn.prompt() conn.setecho(False) return conn
[docs]class TelnetConnection(pxssh.pxssh): # pylint: disable=arguments-differ
[docs] def login(self, server, username, password='', original_prompt=r'[#$]', login_timeout=10, auto_prompt_reset=True, sync_multiplier=1, port=23): cmd = 'telnet -l {} {} {}'.format(username, server, port) spawn._spawn(self, cmd) # pylint: disable=protected-access try: i = self.expect('(?i)(?:password)', timeout=login_timeout) if i == 0: self.sendline(password) i = self.expect([original_prompt, 'Login incorrect'], timeout=login_timeout) if i: raise pxssh.ExceptionPxssh('could not log in: password was incorrect') except TIMEOUT: if not password: # There was no password prompt before TIMEOUT, and we didn't # have a password to enter. Assume everything is OK. pass else: raise pxssh.ExceptionPxssh('could not log in: did not see a password prompt') if not self.sync_original_prompt(sync_multiplier): self.close() raise pxssh.ExceptionPxssh('could not synchronize with original prompt') if auto_prompt_reset: if not self.set_unique_prompt(): self.close() message = 'could not set shell prompt (recieved: {}, expected: {}).' raise pxssh.ExceptionPxssh(message.format(self.before, self.PROMPT)) return True
[docs]def check_keyfile(keyfile): """ keyfile must have the right access premissions in order to be useable. If the specified file doesn't, create a temporary copy and set the right permissions for that. Returns either the ``keyfile`` (if the permissions on it are correct) or the path to a temporary copy with the right permissions. """ desired_mask = stat.S_IWUSR | stat.S_IRUSR actual_mask = os.stat(keyfile).st_mode & 0xFF if actual_mask != desired_mask: tmp_file = os.path.join(tempfile.gettempdir(), os.path.basename(keyfile)) shutil.copy(keyfile, tmp_file) os.chmod(tmp_file, desired_mask) return tmp_file else: # permissions on keyfile are OK return keyfile
[docs]class SshShell(object): default_password_prompt = '[sudo] password' max_cancel_attempts = 5 def __init__(self, password_prompt=None, timeout=10, telnet=False): self.password_prompt = password_prompt if password_prompt is not None else self.default_password_prompt self.timeout = timeout self.telnet = telnet self.conn = None self.lock = threading.Lock() self.connection_lost = False
[docs] def login(self, host, username, password=None, keyfile=None, port=None, timeout=None): # pylint: disable=attribute-defined-outside-init logger.debug('Logging in {}@{}'.format(username, host)) = host self.username = username self.password = password self.keyfile = check_keyfile(keyfile) if keyfile else keyfile self.port = port timeout = self.timeout if timeout is None else timeout self.conn = ssh_get_shell(host, username, password, self.keyfile, port, timeout, self.telnet)
[docs] def push_file(self, source, dest, timeout=30): dest = '{}@{}:{}'.format(self.username,, dest) return self._scp(source, dest, timeout)
[docs] def pull_file(self, source, dest, timeout=30): source = '{}@{}:{}'.format(self.username,, source) return self._scp(source, dest, timeout)
[docs] def background(self, command, stdout=subprocess.PIPE, stderr=subprocess.PIPE): port_string = '-p {}'.format(self.port) if self.port else '' keyfile_string = '-i {}'.format(self.keyfile) if self.keyfile else '' command = '{} {} {} {}@{} {}'.format(ssh, keyfile_string, port_string, self.username,, command) logger.debug(command) if self.password: command = _give_password(self.password, command) return subprocess.Popen(command, stdout=stdout, stderr=stderr, shell=True)
[docs] def reconnect(self): self.conn = ssh_get_shell(, self.username, self.password, self.keyfile, self.port, self.timeout, self.telnet)
[docs] def execute(self, command, timeout=None, check_exit_code=True, as_root=False, strip_colors=True): try: with self.lock: if self.connection_lost: logger.debug('Attempting to reconnect...') self.reconnect() self.connection_lost = False output = self._execute_and_wait_for_prompt(command, timeout, as_root, strip_colors) if check_exit_code: exit_code_text = self._execute_and_wait_for_prompt('echo $?', strip_colors=strip_colors, log=False) try: exit_code = int(exit_code_text.split()[0]) if exit_code: message = 'Got exit code {}\nfrom: {}\nOUTPUT: {}' raise DeviceError(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: self.connection_lost = True raise DeviceError('Connection dropped.')
[docs] def logout(self): logger.debug('Logging out {}@{}'.format(self.username, self.conn.logout()
[docs] def cancel_running_command(self): # simulate impatiently hitting ^C until command prompt appears logger.debug('Sending ^C') for _ in xrange(self.max_cancel_attempts): self.conn.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: command = "sudo -- sh -c '{}'".format(escape_single_quotes(command)) if log: logger.debug(command) self.conn.sendline(command) index = self.conn.expect_exact([self.password_prompt, TIMEOUT], timeout=0.5) if index == 0: self.conn.sendline(self.password) timed_out = self._wait_for_prompt(timeout) output = re.sub(r' \r([^\n])', r'\1', self.conn.before) output = process_backspaces(output) output = re.sub(r'.*?{}'.format(re.escape(command)), '', output, 1).strip() else: if log: logger.debug(command) self.conn.sendline(command) timed_out = self._wait_for_prompt(timeout) # the regex removes line breaks potential introduced when writing # command to shell. output = re.sub(r' \r([^\n])', r'\1', self.conn.before) output = process_backspaces(output) command_index = output.find(command) output = output[command_index + len(command):].strip() if timed_out: self.cancel_running_command() raise TimeoutError(command, output) 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(self.timeout): 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(self.port) if (self.port and self.port != 22) else '' keyfile_string = '-i {}'.format(self.keyfile) if self.keyfile else '' command = '{} -r {} {} {} {}'.format(scp, keyfile_string, port_string, source, dest) pass_string = '' logger.debug(command) if self.password: command = _give_password(self.password, command) try: check_output(command, timeout=timeout, shell=True) except subprocess.CalledProcessError as e: raise CalledProcessErrorWithStderr(e.returncode, e.cmd.replace(pass_string, ''), output=e.output, error=getattr(e, 'error', '')) except TimeoutError as e: raise TimeoutError(e.command.replace(pass_string, ''), e.output)
def _give_password(password, command): if not sshpass: raise HostError('Must have sshpass installed on the host in order to use password-based auth.') pass_string = "sshpass -p '{}' ".format(password) return pass_string + command def _check_env(): global ssh, scp, sshpass # pylint: disable=global-statement if not ssh: ssh = which('ssh') scp = which('scp') sshpass = which('sshpass') if not (ssh and scp): raise HostError('OpenSSH must be installed on the host.')
[docs]def process_backspaces(text): chars = [] for c in text: if c == chr(8) and chars: # backspace chars.pop() else: chars.append(c) return ''.join(chars)