1#
2# Copyright (C) 2013 Intel Corporation
3#
4# SPDX-License-Identifier: MIT
5#
6
7# Provides a class for setting up ssh connections,
8# running commands and copying files to/from a target.
9# It's used by testimage.bbclass and tests in lib/oeqa/runtime.
10
11import subprocess
12import time
13import os
14import select
15
16
17class SSHProcess(object):
18    def __init__(self, **options):
19
20        self.defaultopts = {
21            "stdout": subprocess.PIPE,
22            "stderr": subprocess.STDOUT,
23            "stdin": None,
24            "shell": False,
25            "bufsize": -1,
26            "start_new_session": True,
27        }
28        self.options = dict(self.defaultopts)
29        self.options.update(options)
30        self.status = None
31        self.output = None
32        self.process = None
33        self.starttime = None
34        self.logfile = None
35
36        # Unset DISPLAY which means we won't trigger SSH_ASKPASS
37        env = os.environ.copy()
38        if "DISPLAY" in env:
39            del env['DISPLAY']
40        self.options['env'] = env
41
42    def log(self, msg):
43        if self.logfile:
44            with open(self.logfile, "a") as f:
45               f.write("%s" % msg)
46
47    def _run(self, command, timeout=None, logfile=None):
48        self.logfile = logfile
49        self.starttime = time.time()
50        output = ''
51        self.process = subprocess.Popen(command, **self.options)
52        if timeout:
53            endtime = self.starttime + timeout
54            eof = False
55            while time.time() < endtime and not eof:
56                try:
57                    if select.select([self.process.stdout], [], [], 5)[0] != []:
58                        data = os.read(self.process.stdout.fileno(), 1024)
59                        if not data:
60                            self.process.stdout.close()
61                            eof = True
62                        else:
63                            data = data.decode("utf-8")
64                            output += data
65                            self.log(data)
66                            endtime = time.time() + timeout
67                except InterruptedError:
68                    continue
69
70            # process hasn't returned yet
71            if not eof:
72                self.process.terminate()
73                time.sleep(5)
74                try:
75                    self.process.kill()
76                except OSError:
77                    pass
78                lastline = "\nProcess killed - no output for %d seconds. Total running time: %d seconds." % (timeout, time.time() - self.starttime)
79                self.log(lastline)
80                output += lastline
81        else:
82            output = self.process.communicate()[0]
83            self.log(output.rstrip())
84
85        self.status = self.process.wait()
86        self.output = output.rstrip()
87
88    def run(self, command, timeout=None, logfile=None):
89        try:
90            self._run(command, timeout, logfile)
91        except:
92            # Need to guard against a SystemExit or other exception occuring whilst running
93            # and ensure we don't leave a process behind.
94            if self.process.poll() is None:
95                self.process.kill()
96                self.status = self.process.wait()
97            raise
98        return (self.status, self.output)
99
100class SSHControl(object):
101    def __init__(self, ip, logfile=None, timeout=300, user='root', port=None):
102        self.ip = ip
103        self.defaulttimeout = timeout
104        self.ignore_status = True
105        self.logfile = logfile
106        self.user = user
107        self.ssh_options = [
108                '-o', 'UserKnownHostsFile=/dev/null',
109                '-o', 'StrictHostKeyChecking=no',
110                '-o', 'LogLevel=ERROR'
111                ]
112        self.ssh = ['ssh', '-l', self.user ] + self.ssh_options
113        self.scp = ['scp'] + self.ssh_options
114        if port:
115            self.ssh = self.ssh + [ '-p', port ]
116            self.scp = self.scp + [ '-P', port ]
117
118    def log(self, msg):
119        if self.logfile:
120            with open(self.logfile, "a") as f:
121                f.write("%s\n" % msg)
122
123    def _internal_run(self, command, timeout=None, ignore_status = True):
124        self.log("[Running]$ %s" % " ".join(command))
125
126        proc = SSHProcess()
127        status, output = proc.run(command, timeout, logfile=self.logfile)
128
129        self.log("[Command returned '%d' after %.2f seconds]" % (status, time.time() - proc.starttime))
130
131        if status and not ignore_status:
132            raise AssertionError("Command '%s' returned non-zero exit status %d:\n%s" % (command, status, output))
133
134        return (status, output)
135
136    def run(self, command, timeout=None):
137        """
138        command - ssh command to run
139        timeout=<val> - kill command if there is no output after <val> seconds
140        timeout=None - kill command if there is no output after a default value seconds
141        timeout=0 - no timeout, let command run until it returns
142        """
143
144        command = self.ssh + [self.ip, 'export PATH=/usr/sbin:/sbin:/usr/bin:/bin; ' + command]
145
146        if timeout is None:
147            return self._internal_run(command, self.defaulttimeout, self.ignore_status)
148        if timeout == 0:
149            return self._internal_run(command, None, self.ignore_status)
150        return self._internal_run(command, timeout, self.ignore_status)
151
152    def copy_to(self, localpath, remotepath):
153        if os.path.islink(localpath):
154            localpath = os.path.dirname(localpath) + "/" + os.readlink(localpath)
155        command = self.scp + [localpath, '%s@%s:%s' % (self.user, self.ip, remotepath)]
156        return self._internal_run(command, ignore_status=False)
157
158    def copy_from(self, remotepath, localpath):
159        command = self.scp + ['%s@%s:%s' % (self.user, self.ip, remotepath), localpath]
160        return self._internal_run(command, ignore_status=False)
161
162    def copy_dir_to(self, localpath, remotepath):
163        """
164        Copy recursively localpath directory to remotepath in target.
165        """
166
167        for root, dirs, files in os.walk(localpath):
168            # Create directories in the target as needed
169            for d in dirs:
170                tmp_dir = os.path.join(root, d).replace(localpath, "")
171                new_dir = os.path.join(remotepath, tmp_dir.lstrip("/"))
172                cmd = "mkdir -p %s" % new_dir
173                self.run(cmd)
174
175            # Copy files into the target
176            for f in files:
177                tmp_file = os.path.join(root, f).replace(localpath, "")
178                dst_file = os.path.join(remotepath, tmp_file.lstrip("/"))
179                src_file = os.path.join(root, f)
180                self.copy_to(src_file, dst_file)
181
182
183    def delete_files(self, remotepath, files):
184        """
185        Delete files in target's remote path.
186        """
187
188        cmd = "rm"
189        if not isinstance(files, list):
190            files = [files]
191
192        for f in files:
193            cmd = "%s %s" % (cmd, os.path.join(remotepath, f))
194
195        self.run(cmd)
196
197
198    def delete_dir(self, remotepath):
199        """
200        Delete remotepath directory in target.
201        """
202
203        cmd = "rmdir %s" % remotepath
204        self.run(cmd)
205
206
207    def delete_dir_structure(self, localpath, remotepath):
208        """
209        Delete recursively localpath structure directory in target's remotepath.
210
211        This function is very usefult to delete a package that is installed in
212        the DUT and the host running the test has such package extracted in tmp
213        directory.
214
215        Example:
216            pwd: /home/user/tmp
217            tree:   .
218                    └── work
219                        ├── dir1
220                        │   └── file1
221                        └── dir2
222
223            localpath = "/home/user/tmp" and remotepath = "/home/user"
224
225            With the above variables this function will try to delete the
226            directory in the DUT in this order:
227                /home/user/work/dir1/file1
228                /home/user/work/dir1        (if dir is empty)
229                /home/user/work/dir2        (if dir is empty)
230                /home/user/work             (if dir is empty)
231        """
232
233        for root, dirs, files in os.walk(localpath, topdown=False):
234            # Delete files first
235            tmpdir = os.path.join(root).replace(localpath, "")
236            remotedir = os.path.join(remotepath, tmpdir.lstrip("/"))
237            self.delete_files(remotedir, files)
238
239            # Remove dirs if empty
240            for d in dirs:
241                tmpdir = os.path.join(root, d).replace(localpath, "")
242                remotedir = os.path.join(remotepath, tmpdir.lstrip("/"))
243                self.delete_dir(remotepath)
244