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