| 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.core 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 |
|---|
Chromium Code Reviews
