OLD | NEW |
| (Empty) |
1 # Copyright (c) 2012 The Chromium Authors. All rights reserved. | |
2 # Use of this source code is governed by a BSD-style license that can be | |
3 # found in the LICENSE file. | |
4 """A wrapper around ssh for common operations on a CrOS-based device""" | |
5 import logging | |
6 import os | |
7 import re | |
8 import subprocess | |
9 import sys | |
10 import time | |
11 import tempfile | |
12 | |
13 from telemetry import util | |
14 | |
15 # TODO(nduca): This whole file is built up around making individual ssh calls | |
16 # for each operation. It really could get away with a single ssh session built | |
17 # around pexpect, I suspect, if we wanted it to be faster. But, this was | |
18 # convenient. | |
19 | |
20 def RunCmd(args, cwd=None, quiet=False): | |
21 """Opens a subprocess to execute a program and returns its return value. | |
22 | |
23 Args: | |
24 args: A string or a sequence of program arguments. The program to execute is | |
25 the string or the first item in the args sequence. | |
26 cwd: If not None, the subprocess's current directory will be changed to | |
27 |cwd| before it's executed. | |
28 | |
29 Returns: | |
30 Return code from the command execution. | |
31 """ | |
32 if not quiet: | |
33 logging.debug(' '.join(args) + ' ' + (cwd or '')) | |
34 with open(os.devnull, 'w') as devnull: | |
35 p = subprocess.Popen(args=args, cwd=cwd, stdout=devnull, | |
36 stderr=devnull, stdin=devnull, shell=False) | |
37 return p.wait() | |
38 | |
39 def GetAllCmdOutput(args, cwd=None, quiet=False): | |
40 """Open a subprocess to execute a program and returns its output. | |
41 | |
42 Args: | |
43 args: A string or a sequence of program arguments. The program to execute is | |
44 the string or the first item in the args sequence. | |
45 cwd: If not None, the subprocess's current directory will be changed to | |
46 |cwd| before it's executed. | |
47 | |
48 Returns: | |
49 Captures and returns the command's stdout. | |
50 Prints the command's stderr to logger (which defaults to stdout). | |
51 """ | |
52 if not quiet: | |
53 logging.debug(' '.join(args) + ' ' + (cwd or '')) | |
54 with open(os.devnull, 'w') as devnull: | |
55 p = subprocess.Popen(args=args, cwd=cwd, stdout=subprocess.PIPE, | |
56 stderr=subprocess.PIPE, stdin=devnull, shell=False) | |
57 stdout, stderr = p.communicate() | |
58 if not quiet: | |
59 logging.debug(' > stdout=[%s], stderr=[%s]', stdout, stderr) | |
60 return stdout, stderr | |
61 | |
62 class DeviceSideProcess(object): | |
63 def __init__(self, | |
64 cri, | |
65 device_side_args, | |
66 prevent_output=True, | |
67 extra_ssh_args=None, | |
68 leave_ssh_alive=False, | |
69 env=None, | |
70 login_shell=False): | |
71 | |
72 # Init members first so that Close will always succeed. | |
73 self._cri = cri | |
74 self._proc = None | |
75 self._devnull = open(os.devnull, 'w') | |
76 | |
77 if prevent_output: | |
78 out = self._devnull | |
79 else: | |
80 out = sys.stderr | |
81 | |
82 cri.RmRF('/tmp/cros_interface_remote_device_pid') | |
83 cmd_str = ' '.join(device_side_args) | |
84 if env: | |
85 env_str = ' '.join(['%s=%s' % (k, v) for k, v in env.items()]) | |
86 cmd = env_str + ' ' + cmd_str | |
87 else: | |
88 cmd = cmd_str | |
89 contents = """%s&\n""" % cmd | |
90 contents += 'echo $! > /tmp/cros_interface_remote_device_pid\n' | |
91 cri.PushContents(contents, '/tmp/cros_interface_remote_device_bootstrap.sh') | |
92 | |
93 cmdline = ['/bin/bash'] | |
94 if login_shell: | |
95 cmdline.append('-l') | |
96 cmdline.append('/tmp/cros_interface_remote_device_bootstrap.sh') | |
97 proc = subprocess.Popen( | |
98 cri.FormSSHCommandLine(cmdline, | |
99 extra_ssh_args=extra_ssh_args), | |
100 stdout=out, | |
101 stderr=out, | |
102 stdin=self._devnull, | |
103 shell=False) | |
104 | |
105 time.sleep(0.1) | |
106 def TryGetResult(): | |
107 try: | |
108 self._pid = cri.GetFileContents( | |
109 '/tmp/cros_interface_remote_device_pid').strip() | |
110 return True | |
111 except OSError: | |
112 return False | |
113 try: | |
114 util.WaitFor(TryGetResult, 5) | |
115 except util.TimeoutException: | |
116 raise Exception('Something horrible has happened!') | |
117 | |
118 # Killing the ssh session leaves the process running. We dont | |
119 # need it anymore, unless we have port-forwards. | |
120 if not leave_ssh_alive: | |
121 proc.kill() | |
122 else: | |
123 self._proc = proc | |
124 | |
125 self._pid = int(self._pid) | |
126 if not self.IsAlive(): | |
127 raise OSError('Process did not come up or did not stay alive very long!') | |
128 self._cri = cri | |
129 | |
130 def Close(self, try_sigint_first=False): | |
131 if self.IsAlive(): | |
132 # Try to politely shutdown, first. | |
133 if try_sigint_first: | |
134 logging.debug("kill -INT %i" % self._pid) | |
135 self._cri.GetAllCmdOutput( | |
136 ['kill', '-INT', str(self._pid)], quiet=True) | |
137 try: | |
138 self.Wait(timeout=0.5) | |
139 except util.TimeoutException: | |
140 pass | |
141 | |
142 if self.IsAlive(): | |
143 logging.debug("kill -KILL %i" % self._pid) | |
144 self._cri.GetAllCmdOutput( | |
145 ['kill', '-KILL', str(self._pid)], quiet=True) | |
146 try: | |
147 self.Wait(timeout=5) | |
148 except util.TimeoutException: | |
149 pass | |
150 | |
151 if self.IsAlive(): | |
152 raise Exception('Could not shutdown the process.') | |
153 | |
154 self._cri = None | |
155 if self._proc: | |
156 self._proc.kill() | |
157 self._proc = None | |
158 | |
159 if self._devnull: | |
160 self._devnull.close() | |
161 self._devnull = None | |
162 | |
163 def __enter__(self): | |
164 return self | |
165 | |
166 def __exit__(self, *args): | |
167 self.Close() | |
168 return | |
169 | |
170 def Wait(self, timeout=1): | |
171 if not self._pid: | |
172 raise Exception('Closed') | |
173 def IsDone(): | |
174 return not self.IsAlive() | |
175 util.WaitFor(IsDone, timeout) | |
176 self._pid = None | |
177 | |
178 def IsAlive(self, quiet=True): | |
179 if not self._pid: | |
180 return False | |
181 exists = self._cri.FileExistsOnDevice('/proc/%i/cmdline' % self._pid, | |
182 quiet=quiet) | |
183 return exists | |
184 | |
185 def HasSSH(): | |
186 try: | |
187 RunCmd(['ssh'], quiet=True) | |
188 RunCmd(['scp'], quiet=True) | |
189 logging.debug("HasSSH()->True") | |
190 return True | |
191 except OSError: | |
192 logging.debug("HasSSH()->False") | |
193 return False | |
194 | |
195 class LoginException(Exception): | |
196 pass | |
197 | |
198 class KeylessLoginRequiredException(LoginException): | |
199 pass | |
200 | |
201 class CrOSInterface(object): | |
202 # pylint: disable=R0923 | |
203 def __init__(self, hostname, ssh_identity = None): | |
204 self._hostname = hostname | |
205 self._ssh_identity = None | |
206 self._hostfile = tempfile.NamedTemporaryFile() | |
207 self._hostfile.flush() | |
208 self._ssh_args = ['-o ConnectTimeout=5', | |
209 '-o StrictHostKeyChecking=no', | |
210 '-o KbdInteractiveAuthentication=no', | |
211 '-o PreferredAuthentications=publickey', | |
212 '-o UserKnownHostsFile=%s' % self._hostfile.name] | |
213 | |
214 # List of ports generated from GetRemotePort() that may not be in use yet. | |
215 self._reserved_ports = [] | |
216 | |
217 if ssh_identity: | |
218 self._ssh_identity = os.path.abspath(os.path.expanduser(ssh_identity)) | |
219 | |
220 @property | |
221 def hostname(self): | |
222 return self._hostname | |
223 | |
224 def FormSSHCommandLine(self, args, extra_ssh_args=None): | |
225 full_args = ['ssh', | |
226 '-o ForwardX11=no', | |
227 '-o ForwardX11Trusted=no', | |
228 '-n'] + self._ssh_args | |
229 if self._ssh_identity is not None: | |
230 full_args.extend(['-i', self._ssh_identity]) | |
231 if extra_ssh_args: | |
232 full_args.extend(extra_ssh_args) | |
233 full_args.append('root@%s' % self._hostname) | |
234 full_args.extend(args) | |
235 return full_args | |
236 | |
237 def GetAllCmdOutput(self, args, cwd=None, quiet=False): | |
238 return GetAllCmdOutput(self.FormSSHCommandLine(args), cwd, quiet=quiet) | |
239 | |
240 def _RemoveSSHWarnings(self, toClean): | |
241 """Removes specific ssh warning lines from a string. | |
242 | |
243 Args: | |
244 toClean: A string that may be containing multiple lines. | |
245 | |
246 Returns: | |
247 A copy of toClean with all the Warning lines removed. | |
248 """ | |
249 # Remove the Warning about connecting to a new host for the first time. | |
250 return re.sub('Warning: Permanently added [^\n]* to the list of known ' | |
251 'hosts.\s\n', '', toClean) | |
252 | |
253 def TryLogin(self): | |
254 logging.debug('TryLogin()') | |
255 stdout, stderr = self.GetAllCmdOutput(['echo', '$USER'], quiet=True) | |
256 | |
257 # The initial login will add the host to the hosts file but will also print | |
258 # a warning to stderr that we need to remove. | |
259 stderr = self._RemoveSSHWarnings(stderr) | |
260 if stderr != '': | |
261 if 'Host key verification failed' in stderr: | |
262 raise LoginException(('%s host key verification failed. ' + | |
263 'SSH to it manually to fix connectivity.') % | |
264 self._hostname) | |
265 if 'Operation timed out' in stderr: | |
266 raise LoginException('Timed out while logging into %s' % self._hostname) | |
267 if 'UNPROTECTED PRIVATE KEY FILE!' in stderr: | |
268 raise LoginException('Permissions for %s are too open. To fix this,\n' | |
269 'chmod 600 %s' % (self._ssh_identity, | |
270 self._ssh_identity)) | |
271 if 'Permission denied (publickey,keyboard-interactive)' in stderr: | |
272 raise KeylessLoginRequiredException( | |
273 'Need to set up ssh auth for %s' % self._hostname) | |
274 raise LoginException('While logging into %s, got %s' % ( | |
275 self._hostname, stderr)) | |
276 if stdout != 'root\n': | |
277 raise LoginException( | |
278 'Logged into %s, expected $USER=root, but got %s.' % ( | |
279 self._hostname, stdout)) | |
280 | |
281 def FileExistsOnDevice(self, file_name, quiet=False): | |
282 stdout, stderr = self.GetAllCmdOutput([ | |
283 'if', 'test', '-a', file_name, ';', | |
284 'then', 'echo', '1', ';', | |
285 'fi' | |
286 ], quiet=True) | |
287 if stderr != '': | |
288 if "Connection timed out" in stderr: | |
289 raise OSError('Machine wasn\'t responding to ssh: %s' % | |
290 stderr) | |
291 raise OSError('Unepected error: %s' % stderr) | |
292 exists = stdout == '1\n' | |
293 if not quiet: | |
294 logging.debug("FileExistsOnDevice(<text>, %s)->%s" % ( | |
295 file_name, exists)) | |
296 return exists | |
297 | |
298 def PushFile(self, filename, remote_filename): | |
299 args = ['scp', '-r' ] + self._ssh_args | |
300 if self._ssh_identity: | |
301 args.extend(['-i', self._ssh_identity]) | |
302 | |
303 args.extend([os.path.abspath(filename), | |
304 'root@%s:%s' % (self._hostname, remote_filename)]) | |
305 | |
306 stdout, stderr = GetAllCmdOutput(args, quiet=True) | |
307 if stderr != '': | |
308 assert 'No such file or directory' in stderr | |
309 raise OSError | |
310 | |
311 def PushContents(self, text, remote_filename): | |
312 logging.debug("PushContents(<text>, %s)" % remote_filename) | |
313 with tempfile.NamedTemporaryFile() as f: | |
314 f.write(text) | |
315 f.flush() | |
316 self.PushFile(f.name, remote_filename) | |
317 | |
318 def GetFileContents(self, filename): | |
319 with tempfile.NamedTemporaryFile() as f: | |
320 args = ['scp'] + self._ssh_args | |
321 if self._ssh_identity: | |
322 args.extend(['-i', self._ssh_identity]) | |
323 | |
324 args.extend(['root@%s:%s' % (self._hostname, filename), | |
325 os.path.abspath(f.name)]) | |
326 | |
327 stdout, stderr = GetAllCmdOutput(args, quiet=True) | |
328 | |
329 if stderr != '': | |
330 assert 'No such file or directory' in stderr | |
331 raise OSError | |
332 | |
333 with open(f.name, 'r') as f2: | |
334 res = f2.read() | |
335 logging.debug("GetFileContents(%s)->%s" % (filename, res)) | |
336 return res | |
337 | |
338 def ListProcesses(self): | |
339 stdout, stderr = self.GetAllCmdOutput([ | |
340 '/bin/ps', '--no-headers', | |
341 '-A', | |
342 '-o', 'pid,args'], quiet=True) | |
343 assert stderr == '' | |
344 procs = [] | |
345 for l in stdout.split('\n'): # pylint: disable=E1103 | |
346 if l == '': | |
347 continue | |
348 m = re.match('^\s*(\d+)\s+(.+)', l, re.DOTALL) | |
349 assert m | |
350 procs.append(m.groups()) | |
351 logging.debug("ListProcesses(<predicate>)->[%i processes]" % len(procs)) | |
352 return procs | |
353 | |
354 def RmRF(self, filename): | |
355 logging.debug("rm -rf %s" % filename) | |
356 self.GetCmdOutput(['rm', '-rf', filename], quiet=True) | |
357 | |
358 def KillAllMatching(self, predicate): | |
359 kills = ['kill', '-KILL'] | |
360 for p in self.ListProcesses(): | |
361 if predicate(p[1]): | |
362 logging.info('Killing %s', repr(p)) | |
363 kills.append(p[0]) | |
364 logging.debug("KillAllMatching(<predicate>)->%i" % (len(kills) - 2)) | |
365 if len(kills) > 2: | |
366 self.GetCmdOutput(kills, quiet=True) | |
367 return len(kills) - 2 | |
368 | |
369 def IsServiceRunning(self, service_name): | |
370 stdout, stderr = self.GetAllCmdOutput([ | |
371 'status', service_name], quiet=True) | |
372 assert stderr == '' | |
373 running = 'running, process' in stdout | |
374 logging.debug("IsServiceRunning(%s)->%s" % (service_name, running)) | |
375 return running | |
376 | |
377 def GetCmdOutput(self, args, quiet=False): | |
378 stdout, stderr = self.GetAllCmdOutput(args, quiet=True) | |
379 assert stderr == '' | |
380 if not quiet: | |
381 logging.debug("GetCmdOutput(%s)->%s" % (repr(args), stdout)) | |
382 return stdout | |
383 | |
384 def GetRemotePort(self): | |
385 netstat = self.GetAllCmdOutput(['netstat', '-ant']) | |
386 netstat = netstat[0].split('\n') | |
387 ports_in_use = [] | |
388 | |
389 for line in netstat[2:]: | |
390 if not line: | |
391 continue | |
392 address_in_use = line.split()[3] | |
393 port_in_use = address_in_use.split(':')[-1] | |
394 ports_in_use.append(int(port_in_use)) | |
395 | |
396 ports_in_use.extend(self._reserved_ports) | |
397 | |
398 new_port = sorted(ports_in_use)[-1] + 1 | |
399 self._reserved_ports.append(new_port) | |
400 | |
401 return new_port | |
402 | |
403 def IsHTTPServerRunningOnPort(self, port): | |
404 wget_output = self.GetAllCmdOutput( | |
405 ['wget', 'localhost:%i' % (port), '-T1', '-t1']) | |
406 | |
407 if 'Connection refused' in wget_output[1]: | |
408 return False | |
409 | |
410 return True | |
OLD | NEW |