blob: 49a07264c67959e080d667b67bffb71a808c7137 [file] [log] [blame]
Brad Bishopc342db32019-05-15 21:57:59 -04001#
Patrick Williamsc124f4f2015-09-15 14:41:29 -05002# Copyright (C) 2013 Intel Corporation
3#
Brad Bishopc342db32019-05-15 21:57:59 -04004# SPDX-License-Identifier: MIT
5#
Patrick Williamsc124f4f2015-09-15 14:41:29 -05006
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 "preexec_fn": os.setsid,
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:
Patrick Williamsc0f7c042017-02-23 20:41:17 -060056 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
Patrick Williamsc124f4f2015-09-15 14:41:29 -050069
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
Patrick Williamsd8c66bc2016-06-20 12:57:21 -0500144 command = self.ssh + [self.ip, 'export PATH=/usr/sbin:/sbin:/usr/bin:/bin; ' + command]
Patrick Williamsc124f4f2015-09-15 14:41:29 -0500145
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):
Patrick Williamsc0f7c042017-02-23 20:41:17 -0600153 if os.path.islink(localpath):
Brad Bishopd7bf8c12018-02-25 22:55:05 -0500154 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)
Patrick Williamsc124f4f2015-09-15 14:41:29 -0500157
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)
Patrick Williamsc0f7c042017-02-23 20:41:17 -0600161
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)