blob: 71744050733142d28176fbdf9acd02fbfce9b0e1 [file] [log] [blame]
#!/usr/bin/env python
# Copyright 2018 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Implementation of using SSH."""
import os
import pipes
import platform
import shlex
import subprocess
import threading
import time
import six
from six.moves import queue
import graphyte_common # pylint: disable=unused-import
from graphyte import link
from graphyte.default_setting import logger
from graphyte.utils.file_utils import UnopenedTemporaryFile
from graphyte.utils import graphyte_utils
from graphyte.utils import process_utils
from graphyte.utils import sync_utils
from graphyte.utils import type_utils
class ClientNotExistError(Exception):
def __str__(self):
return 'There is no DHCP client registered.'
class SSHLink(link.DeviceLink):
"""A DUT target that is connected via SSH interface.
host: A string for SSH host IP.
user: A string for the user accont to login. Defaults to 'root'.
port: An integer for the SSH port on remote host.
identify: An identity file to specify credential.
connect_timeout: An integer for ssh(1) connection timeout in seconds.
control_persist: An integer for ssh(1) to keep master connection remain
opened for given seconds, or None to not using master control.
link_options example:
'link_class': 'SSHLink',
'host': '',
'identity': '/path/to/identity/file'
def __init__(self, host, user='root', port=22, identity=None,
connect_timeout=5, control_persist=300): = host
self.user = user
self.port = port
self.identity = identity
self.connect_timeout = connect_timeout
self.control_persist = control_persist
if self.identity is not None:
self.identity = graphyte_utils.SearchConfig(self.identity)
os.chmod(self.identity, 0o600)
except Exception:
logger.error('SSHLink: Failed to locate identity file "%s"',
self.identity = None
def __str__(self):
return "SSH to %s@%s:%s with identity='%s'" % (
self.user,, self.port, self.identity)
def _Signature(self, is_scp=False):
"""Generates the ssh command signature.
is_scp: A boolean flag indicating if the signature is made for scp.
A pair of signature in (sig, options). The 'sig' is a string representing
remote ssh user and host. 'options' is a list of required command line
if self.user:
sig = '%s@%s' % (self.user,
sig =
options = ['-o', 'UserKnownHostsFile=/dev/null',
'-o', 'StrictHostKeyChecking=no',
'-o', 'ConnectTimeout=%d' % self.connect_timeout]
if self.control_persist is not None:
options += ['-o', 'ControlMaster=auto',
'-o', 'ControlPath=/tmp/.ssh-%r@%h:%p',
'-o', 'ControlPersist=%s' % self.control_persist]
if self.port:
options += ['-P' if is_scp else '-p', str(self.port)]
if self.identity:
options += ['-i', self.identity]
return sig, options
def _ToUnixPath(self, path):
if platform.system() != 'Windows':
return path # assume that we are in Unix machine already
path = os.path.abspath(path)
drive, path = os.path.splitdrive(path)
return '/' + drive[0] + path.replace(os.path.sep, '/')
def Push(self, local, remote):
"""See DUTLink.Push"""
remote_sig, options = self._Signature(True)
local = self._ToUnixPath(local)
return subprocess.check_call(['scp'] + options +
[local, '%s:%s' % (remote_sig, remote)])
def Pull(self, remote, local=None):
"""See DUTLink.Pull"""
if local is None:
with UnopenedTemporaryFile() as path:
self.Pull(remote, path)
with open(path) as f:
local = self._ToUnixPath(local)
remote_sig, options = self._Signature(True)
subprocess.check_call(['scp'] + options +
['%s:%s' % (remote_sig, remote), local])
def _StartWatcher(self, subproc):
watcher = self.__class__.ControlMasterWatcher(self)
watcher.Start() # make sure the watcher is running
watcher.AddProcess(, os.getpid())
def Shell(self, command, stdin=None, stdout=None, stderr=None):
"""See DUTLink.Shell"""
remote_sig, options = self._Signature(False)
if isinstance(command, six.string_types):
# since the shell command is passed through SSH client,
# the entire command is argument for SSH command
# e.g. command = 'VAR=xxx cmd arg0 arg1...'
# will become: 'ssh' <SSH OPTIONS> 'VAR=xxx' 'cmd' 'arg0' 'arg1' ...
command = shlex.split(command)
command = [pipes.quote(c) for c in command]
command = ['ssh'] + options + [remote_sig] + command
shell = False
logger.debug('SSHLink: Run [%r]', command)
close_fds = True
if (platform.system() == 'Windows' and
(stdin, stdout, stderr) != (None, None, None)):
close_fds = False
proc = subprocess.Popen(command, shell=shell, close_fds=close_fds,
stdout=stdout, stderr=stderr)
return proc
def IsReady(self):
"""See DUTLink.IsReady"""
return['ping', '-w', '1', '-c', '1',]) == 0
except ClientNotExistError:
return False
class ControlMasterWatcher(object):
__metaclass__ = type_utils.Singleton
def __init__(self, link_instance):
assert isinstance(link_instance, SSHLink)
self._link = link_instance
self._thread = threading.Thread(target=self.Run)
self._proc_queue = queue.Queue()
self._user = self._link.user
self._host =
self._port = self._link.port
self._link_class_name = self._link.__class__.__name__
def IsRunning(self):
if not self._thread:
return False
if not self._thread.is_alive():
self._thread = None
return False
return True
def Start(self):
if self.IsRunning():
if self._link.control_persist is None:
logger.debug('%s %s@%s:%s is not using control master, don\'t start',
self._link_class_name, self._user, self._host, self._port)
self._thread = process_utils.StartDaemonThread(target=self.Run)
def AddProcess(self, pid, ppid=None):
"""Add an SSH process to monitor.
If any of added SSH process is still running, ControlMasterWatcher will
keep monitoring network connectivity. If network is down, control master
will be killed.
pid: PID of process using SSH
ppid: parent PID of given process
if not self.IsRunning():
logger.warning('Watcher is not running, %d is not added', pid)
self._proc_queue.put((pid, ppid))
def Run(self):'start monitoring control master')
# an alias to prevent duplicated pylint warnings
# pylint: disable=protected-access
_GetLinkSignature = self._link._Signature
def _IsControlMasterRunning():
sig, options = _GetLinkSignature(False)
['ssh', '-O', 'check'] + options + [sig, 'true'],
stderr=subprocess.PIPE) == 0
def _StopControlMaster():
sig, options = _GetLinkSignature(False)['ssh', '-O', 'exit'] + options + [sig, 'true'])
def _CallTrue():
sig, options = _GetLinkSignature(False)
proc = subprocess.Popen(['ssh'] + options + [sig, 'true'])
returncode = proc.poll()
if returncode != 0:
return False
return True
def _PollingCallback(is_process_alive):
if not is_process_alive:
return True # returns True to stop polling
if not _IsControlMasterRunning():'control master is not running, skipped')
if not _CallTrue():'loss connection, stopping control master')
except Exception:'monitoring %s to %s@%s:%s',
self._user, self._host, self._port, exc_info=True)
while True:
# get a new process from queue to monitor
# since Queue.get will block if queue is empty, we don't need to sleep
pid, ppid = self._proc_queue.get()
logger.debug('start monitoring control master until %d terminates', pid)
lambda: process_utils.IsProcessAlive(pid, ppid),