| # -*- coding: utf-8 -*- |
| # Copyright (C) 2013 Intel Corporation |
| # |
| # Released under the MIT license (see COPYING.MIT) |
| |
| # Provides a class for setting up ssh connections, |
| # running commands and copying files to/from a target. |
| # It's used by testimage.bbclass and tests in lib/oeqa/runtime. |
| |
| import subprocess |
| import time |
| import os |
| import select |
| |
| |
| class SSHProcess(object): |
| def __init__(self, **options): |
| |
| self.defaultopts = { |
| "stdout": subprocess.PIPE, |
| "stderr": subprocess.STDOUT, |
| "stdin": None, |
| "shell": False, |
| "bufsize": -1, |
| "preexec_fn": os.setsid, |
| } |
| self.options = dict(self.defaultopts) |
| self.options.update(options) |
| self.status = None |
| self.output = None |
| self.process = None |
| self.starttime = None |
| self.logfile = None |
| |
| # Unset DISPLAY which means we won't trigger SSH_ASKPASS |
| env = os.environ.copy() |
| if "DISPLAY" in env: |
| del env['DISPLAY'] |
| self.options['env'] = env |
| |
| def log(self, msg): |
| if self.logfile: |
| with open(self.logfile, "a") as f: |
| f.write("%s" % msg) |
| |
| def _run(self, command, timeout=None, logfile=None): |
| self.logfile = logfile |
| self.starttime = time.time() |
| output = '' |
| self.process = subprocess.Popen(command, **self.options) |
| if timeout: |
| endtime = self.starttime + timeout |
| eof = False |
| while time.time() < endtime and not eof: |
| try: |
| if select.select([self.process.stdout], [], [], 5)[0] != []: |
| data = os.read(self.process.stdout.fileno(), 1024) |
| if not data: |
| self.process.stdout.close() |
| eof = True |
| else: |
| data = data.decode("utf-8") |
| output += data |
| self.log(data) |
| endtime = time.time() + timeout |
| except InterruptedError: |
| continue |
| |
| # process hasn't returned yet |
| if not eof: |
| self.process.terminate() |
| time.sleep(5) |
| try: |
| self.process.kill() |
| except OSError: |
| pass |
| lastline = "\nProcess killed - no output for %d seconds. Total running time: %d seconds." % (timeout, time.time() - self.starttime) |
| self.log(lastline) |
| output += lastline |
| else: |
| output = self.process.communicate()[0] |
| self.log(output.rstrip()) |
| |
| self.status = self.process.wait() |
| self.output = output.rstrip() |
| |
| def run(self, command, timeout=None, logfile=None): |
| try: |
| self._run(command, timeout, logfile) |
| except: |
| # Need to guard against a SystemExit or other exception occuring whilst running |
| # and ensure we don't leave a process behind. |
| if self.process.poll() is None: |
| self.process.kill() |
| self.status = self.process.wait() |
| raise |
| return (self.status, self.output) |
| |
| class SSHControl(object): |
| def __init__(self, ip, logfile=None, timeout=300, user='root', port=None): |
| self.ip = ip |
| self.defaulttimeout = timeout |
| self.ignore_status = True |
| self.logfile = logfile |
| self.user = user |
| self.ssh_options = [ |
| '-o', 'UserKnownHostsFile=/dev/null', |
| '-o', 'StrictHostKeyChecking=no', |
| '-o', 'LogLevel=ERROR' |
| ] |
| self.ssh = ['ssh', '-l', self.user ] + self.ssh_options |
| self.scp = ['scp'] + self.ssh_options |
| if port: |
| self.ssh = self.ssh + [ '-p', port ] |
| self.scp = self.scp + [ '-P', port ] |
| |
| def log(self, msg): |
| if self.logfile: |
| with open(self.logfile, "a") as f: |
| f.write("%s\n" % msg) |
| |
| def _internal_run(self, command, timeout=None, ignore_status = True): |
| self.log("[Running]$ %s" % " ".join(command)) |
| |
| proc = SSHProcess() |
| status, output = proc.run(command, timeout, logfile=self.logfile) |
| |
| self.log("[Command returned '%d' after %.2f seconds]" % (status, time.time() - proc.starttime)) |
| |
| if status and not ignore_status: |
| raise AssertionError("Command '%s' returned non-zero exit status %d:\n%s" % (command, status, output)) |
| |
| return (status, output) |
| |
| def run(self, command, timeout=None): |
| """ |
| command - ssh command to run |
| timeout=<val> - kill command if there is no output after <val> seconds |
| timeout=None - kill command if there is no output after a default value seconds |
| timeout=0 - no timeout, let command run until it returns |
| """ |
| |
| command = self.ssh + [self.ip, 'export PATH=/usr/sbin:/sbin:/usr/bin:/bin; ' + command] |
| |
| if timeout is None: |
| return self._internal_run(command, self.defaulttimeout, self.ignore_status) |
| if timeout == 0: |
| return self._internal_run(command, None, self.ignore_status) |
| return self._internal_run(command, timeout, self.ignore_status) |
| |
| def copy_to(self, localpath, remotepath): |
| if os.path.islink(localpath): |
| localpath = os.path.dirname(localpath) + "/" + os.readlink(localpath) |
| command = self.scp + [localpath, '%s@%s:%s' % (self.user, self.ip, remotepath)] |
| return self._internal_run(command, ignore_status=False) |
| |
| def copy_from(self, remotepath, localpath): |
| command = self.scp + ['%s@%s:%s' % (self.user, self.ip, remotepath), localpath] |
| return self._internal_run(command, ignore_status=False) |
| |
| def copy_dir_to(self, localpath, remotepath): |
| """ |
| Copy recursively localpath directory to remotepath in target. |
| """ |
| |
| for root, dirs, files in os.walk(localpath): |
| # Create directories in the target as needed |
| for d in dirs: |
| tmp_dir = os.path.join(root, d).replace(localpath, "") |
| new_dir = os.path.join(remotepath, tmp_dir.lstrip("/")) |
| cmd = "mkdir -p %s" % new_dir |
| self.run(cmd) |
| |
| # Copy files into the target |
| for f in files: |
| tmp_file = os.path.join(root, f).replace(localpath, "") |
| dst_file = os.path.join(remotepath, tmp_file.lstrip("/")) |
| src_file = os.path.join(root, f) |
| self.copy_to(src_file, dst_file) |
| |
| |
| def delete_files(self, remotepath, files): |
| """ |
| Delete files in target's remote path. |
| """ |
| |
| cmd = "rm" |
| if not isinstance(files, list): |
| files = [files] |
| |
| for f in files: |
| cmd = "%s %s" % (cmd, os.path.join(remotepath, f)) |
| |
| self.run(cmd) |
| |
| |
| def delete_dir(self, remotepath): |
| """ |
| Delete remotepath directory in target. |
| """ |
| |
| cmd = "rmdir %s" % remotepath |
| self.run(cmd) |
| |
| |
| def delete_dir_structure(self, localpath, remotepath): |
| """ |
| Delete recursively localpath structure directory in target's remotepath. |
| |
| This function is very usefult to delete a package that is installed in |
| the DUT and the host running the test has such package extracted in tmp |
| directory. |
| |
| Example: |
| pwd: /home/user/tmp |
| tree: . |
| └── work |
| ├── dir1 |
| │ └── file1 |
| └── dir2 |
| |
| localpath = "/home/user/tmp" and remotepath = "/home/user" |
| |
| With the above variables this function will try to delete the |
| directory in the DUT in this order: |
| /home/user/work/dir1/file1 |
| /home/user/work/dir1 (if dir is empty) |
| /home/user/work/dir2 (if dir is empty) |
| /home/user/work (if dir is empty) |
| """ |
| |
| for root, dirs, files in os.walk(localpath, topdown=False): |
| # Delete files first |
| tmpdir = os.path.join(root).replace(localpath, "") |
| remotedir = os.path.join(remotepath, tmpdir.lstrip("/")) |
| self.delete_files(remotedir, files) |
| |
| # Remove dirs if empty |
| for d in dirs: |
| tmpdir = os.path.join(root, d).replace(localpath, "") |
| remotedir = os.path.join(remotepath, tmpdir.lstrip("/")) |
| self.delete_dir(remotepath) |