1# Copyright 2011 The Chromium OS 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"""Utilities to run commands in outside/inside chroot and on the board.""" 5 6from __future__ import print_function 7 8import getpass 9import os 10import re 11import select 12import signal 13import subprocess 14import sys 15import tempfile 16import time 17 18import logger 19import misc 20 21mock_default = False 22 23LOG_LEVEL = ('none', 'quiet', 'average', 'verbose') 24 25 26def InitCommandExecuter(mock=False): 27 # pylint: disable=global-statement 28 global mock_default 29 # Whether to default to a mock command executer or not 30 mock_default = mock 31 32 33def GetCommandExecuter(logger_to_set=None, mock=False, log_level='verbose'): 34 # If the default is a mock executer, always return one. 35 if mock_default or mock: 36 return MockCommandExecuter(log_level, logger_to_set) 37 else: 38 return CommandExecuter(log_level, logger_to_set) 39 40 41class CommandExecuter(object): 42 """Provides several methods to execute commands on several environments.""" 43 44 def __init__(self, log_level, logger_to_set=None): 45 self.log_level = log_level 46 if log_level == 'none': 47 self.logger = None 48 else: 49 if logger_to_set is not None: 50 self.logger = logger_to_set 51 else: 52 self.logger = logger.GetLogger() 53 54 def GetLogLevel(self): 55 return self.log_level 56 57 def SetLogLevel(self, log_level): 58 self.log_level = log_level 59 60 def RunCommandGeneric(self, 61 cmd, 62 return_output=False, 63 machine=None, 64 username=None, 65 command_terminator=None, 66 command_timeout=None, 67 terminated_timeout=10, 68 print_to_console=True, 69 except_handler=lambda p, e: None): 70 """Run a command. 71 72 Returns triplet (returncode, stdout, stderr). 73 """ 74 75 cmd = str(cmd) 76 77 if self.log_level == 'quiet': 78 print_to_console = False 79 80 if self.log_level == 'verbose': 81 self.logger.LogCmd(cmd, machine, username, print_to_console) 82 elif self.logger: 83 self.logger.LogCmdToFileOnly(cmd, machine, username) 84 if command_terminator and command_terminator.IsTerminated(): 85 if self.logger: 86 self.logger.LogError('Command was terminated!', print_to_console) 87 return (1, '', '') 88 89 if machine is not None: 90 user = '' 91 if username is not None: 92 user = username + '@' 93 cmd = "ssh -t -t %s%s -- '%s'" % (user, machine, cmd) 94 95 # We use setsid so that the child will have a different session id 96 # and we can easily kill the process group. This is also important 97 # because the child will be disassociated from the parent terminal. 98 # In this way the child cannot mess the parent's terminal. 99 p = None 100 try: 101 p = subprocess.Popen(cmd, 102 stdout=subprocess.PIPE, 103 stderr=subprocess.PIPE, 104 shell=True, 105 preexec_fn=os.setsid) 106 107 full_stdout = '' 108 full_stderr = '' 109 110 # Pull output from pipes, send it to file/stdout/string 111 out = err = None 112 pipes = [p.stdout, p.stderr] 113 114 my_poll = select.poll() 115 my_poll.register(p.stdout, select.POLLIN) 116 my_poll.register(p.stderr, select.POLLIN) 117 118 terminated_time = None 119 started_time = time.time() 120 121 while len(pipes): 122 if command_terminator and command_terminator.IsTerminated(): 123 os.killpg(os.getpgid(p.pid), signal.SIGTERM) 124 if self.logger: 125 self.logger.LogError('Command received termination request. ' 126 'Killed child process group.', 127 print_to_console) 128 break 129 130 l = my_poll.poll(100) 131 for (fd, _) in l: 132 if fd == p.stdout.fileno(): 133 out = os.read(p.stdout.fileno(), 16384) 134 if return_output: 135 full_stdout += out 136 if self.logger: 137 self.logger.LogCommandOutput(out, print_to_console) 138 if out == '': 139 pipes.remove(p.stdout) 140 my_poll.unregister(p.stdout) 141 if fd == p.stderr.fileno(): 142 err = os.read(p.stderr.fileno(), 16384) 143 if return_output: 144 full_stderr += err 145 if self.logger: 146 self.logger.LogCommandError(err, print_to_console) 147 if err == '': 148 pipes.remove(p.stderr) 149 my_poll.unregister(p.stderr) 150 151 if p.poll() is not None: 152 if terminated_time is None: 153 terminated_time = time.time() 154 elif (terminated_timeout is not None and 155 time.time() - terminated_time > terminated_timeout): 156 if self.logger: 157 self.logger.LogWarning('Timeout of %s seconds reached since ' 158 'process termination.' % 159 terminated_timeout, 160 print_to_console) 161 break 162 163 if (command_timeout is not None and 164 time.time() - started_time > command_timeout): 165 os.killpg(os.getpgid(p.pid), signal.SIGTERM) 166 if self.logger: 167 self.logger.LogWarning('Timeout of %s seconds reached since process' 168 'started. Killed child process group.' % 169 command_timeout, print_to_console) 170 break 171 172 if out == err == '': 173 break 174 175 p.wait() 176 if return_output: 177 return (p.returncode, full_stdout, full_stderr) 178 return (p.returncode, '', '') 179 except BaseException as e: 180 except_handler(p, e) 181 raise 182 183 def RunCommand(self, *args, **kwargs): 184 """Run a command. 185 186 Takes the same arguments as RunCommandGeneric except for return_output. 187 Returns a single value returncode. 188 """ 189 # Make sure that args does not overwrite 'return_output' 190 assert len(args) <= 1 191 assert 'return_output' not in kwargs 192 kwargs['return_output'] = False 193 return self.RunCommandGeneric(*args, **kwargs)[0] 194 195 def RunCommandWExceptionCleanup(self, *args, **kwargs): 196 """Run a command and kill process if exception is thrown. 197 198 Takes the same arguments as RunCommandGeneric except for except_handler. 199 Returns same as RunCommandGeneric. 200 """ 201 202 def KillProc(proc, _): 203 if proc: 204 os.killpg(os.getpgid(proc.pid), signal.SIGTERM) 205 206 # Make sure that args does not overwrite 'except_handler' 207 assert len(args) <= 8 208 assert 'except_handler' not in kwargs 209 kwargs['except_handler'] = KillProc 210 return self.RunCommandGeneric(*args, **kwargs) 211 212 def RunCommandWOutput(self, *args, **kwargs): 213 """Run a command. 214 215 Takes the same arguments as RunCommandGeneric except for return_output. 216 Returns a triplet (returncode, stdout, stderr). 217 """ 218 # Make sure that args does not overwrite 'return_output' 219 assert len(args) <= 1 220 assert 'return_output' not in kwargs 221 kwargs['return_output'] = True 222 return self.RunCommandGeneric(*args, **kwargs) 223 224 def RemoteAccessInitCommand(self, chromeos_root, machine): 225 command = '' 226 command += '\nset -- --remote=' + machine 227 command += '\n. ' + chromeos_root + '/src/scripts/common.sh' 228 command += '\n. ' + chromeos_root + '/src/scripts/remote_access.sh' 229 command += '\nTMP=$(mktemp -d)' 230 command += "\nFLAGS \"$@\" || exit 1" 231 command += '\nremote_access_init' 232 return command 233 234 def WriteToTempShFile(self, contents): 235 handle, command_file = tempfile.mkstemp(prefix=os.uname()[1], suffix='.sh') 236 os.write(handle, '#!/bin/bash\n') 237 os.write(handle, contents) 238 os.close(handle) 239 return command_file 240 241 def CrosLearnBoard(self, chromeos_root, machine): 242 command = self.RemoteAccessInitCommand(chromeos_root, machine) 243 command += '\nlearn_board' 244 command += '\necho ${FLAGS_board}' 245 retval, output, _ = self.RunCommandWOutput(command) 246 if self.logger: 247 self.logger.LogFatalIf(retval, 'learn_board command failed') 248 elif retval: 249 sys.exit(1) 250 return output.split()[-1] 251 252 def CrosRunCommandGeneric(self, 253 cmd, 254 return_output=False, 255 machine=None, 256 command_terminator=None, 257 chromeos_root=None, 258 command_timeout=None, 259 terminated_timeout=10, 260 print_to_console=True): 261 """Run a command on a ChromeOS box. 262 263 Returns triplet (returncode, stdout, stderr). 264 """ 265 266 if self.log_level != 'verbose': 267 print_to_console = False 268 269 if self.logger: 270 self.logger.LogCmd(cmd, print_to_console=print_to_console) 271 self.logger.LogFatalIf(not machine, 'No machine provided!') 272 self.logger.LogFatalIf(not chromeos_root, 'chromeos_root not given!') 273 else: 274 if not chromeos_root or not machine: 275 sys.exit(1) 276 chromeos_root = os.path.expanduser(chromeos_root) 277 278 # Write all commands to a file. 279 command_file = self.WriteToTempShFile(cmd) 280 retval = self.CopyFiles(command_file, 281 command_file, 282 dest_machine=machine, 283 command_terminator=command_terminator, 284 chromeos_root=chromeos_root, 285 dest_cros=True, 286 recursive=False, 287 print_to_console=print_to_console) 288 if retval: 289 if self.logger: 290 self.logger.LogError('Could not run remote command on machine.' 291 ' Is the machine up?') 292 return (retval, '', '') 293 294 command = self.RemoteAccessInitCommand(chromeos_root, machine) 295 command += '\nremote_sh bash %s' % command_file 296 command += "\nl_retval=$?; echo \"$REMOTE_OUT\"; exit $l_retval" 297 retval = self.RunCommandGeneric(command, 298 return_output, 299 command_terminator=command_terminator, 300 command_timeout=command_timeout, 301 terminated_timeout=terminated_timeout, 302 print_to_console=print_to_console) 303 if return_output: 304 connect_signature = ( 305 'Initiating first contact with remote host\n' + 'Connection OK\n') 306 connect_signature_re = re.compile(connect_signature) 307 modded_retval = list(retval) 308 modded_retval[1] = connect_signature_re.sub('', retval[1]) 309 return modded_retval 310 return retval 311 312 def CrosRunCommand(self, *args, **kwargs): 313 """Run a command on a ChromeOS box. 314 315 Takes the same arguments as CrosRunCommandGeneric except for return_output. 316 Returns a single value returncode. 317 """ 318 # Make sure that args does not overwrite 'return_output' 319 assert len(args) <= 1 320 assert 'return_output' not in kwargs 321 kwargs['return_output'] = False 322 return self.CrosRunCommandGeneric(*args, **kwargs)[0] 323 324 def CrosRunCommandWOutput(self, *args, **kwargs): 325 """Run a command on a ChromeOS box. 326 327 Takes the same arguments as CrosRunCommandGeneric except for return_output. 328 Returns a triplet (returncode, stdout, stderr). 329 """ 330 # Make sure that args does not overwrite 'return_output' 331 assert len(args) <= 1 332 assert 'return_output' not in kwargs 333 kwargs['return_output'] = True 334 return self.CrosRunCommandGeneric(*args, **kwargs) 335 336 def ChrootRunCommandGeneric(self, 337 chromeos_root, 338 command, 339 return_output=False, 340 command_terminator=None, 341 command_timeout=None, 342 terminated_timeout=10, 343 print_to_console=True, 344 cros_sdk_options=''): 345 """Runs a command within the chroot. 346 347 Returns triplet (returncode, stdout, stderr). 348 """ 349 350 if self.log_level != 'verbose': 351 print_to_console = False 352 353 if self.logger: 354 self.logger.LogCmd(command, print_to_console=print_to_console) 355 356 handle, command_file = tempfile.mkstemp( 357 dir=os.path.join(chromeos_root, 'src/scripts'), 358 suffix='.sh', 359 prefix='in_chroot_cmd') 360 os.write(handle, '#!/bin/bash\n') 361 os.write(handle, command) 362 os.write(handle, '\n') 363 os.close(handle) 364 365 os.chmod(command_file, 0777) 366 367 # if return_output is set, run a dummy command first to make sure that 368 # the chroot already exists. We want the final returned output to skip 369 # the output from chroot creation steps. 370 if return_output: 371 ret = self.RunCommand('cd %s; cros_sdk %s -- true' % 372 (chromeos_root, cros_sdk_options)) 373 if ret: 374 return (ret, '', '') 375 376 # Run command_file inside the chroot, making sure that any "~" is expanded 377 # by the shell inside the chroot, not outside. 378 command = ("cd %s; cros_sdk %s -- bash -c '%s/%s'" % 379 (chromeos_root, cros_sdk_options, misc.CHROMEOS_SCRIPTS_DIR, 380 os.path.basename(command_file))) 381 ret = self.RunCommandGeneric(command, 382 return_output, 383 command_terminator=command_terminator, 384 command_timeout=command_timeout, 385 terminated_timeout=terminated_timeout, 386 print_to_console=print_to_console) 387 os.remove(command_file) 388 return ret 389 390 def ChrootRunCommand(self, *args, **kwargs): 391 """Runs a command within the chroot. 392 393 Takes the same arguments as ChrootRunCommandGeneric except for 394 return_output. 395 Returns a single value returncode. 396 """ 397 # Make sure that args does not overwrite 'return_output' 398 assert len(args) <= 2 399 assert 'return_output' not in kwargs 400 kwargs['return_output'] = False 401 return self.ChrootRunCommandGeneric(*args, **kwargs)[0] 402 403 def ChrootRunCommandWOutput(self, *args, **kwargs): 404 """Runs a command within the chroot. 405 406 Takes the same arguments as ChrootRunCommandGeneric except for 407 return_output. 408 Returns a triplet (returncode, stdout, stderr). 409 """ 410 # Make sure that args does not overwrite 'return_output' 411 assert len(args) <= 2 412 assert 'return_output' not in kwargs 413 kwargs['return_output'] = True 414 return self.ChrootRunCommandGeneric(*args, **kwargs) 415 416 def RunCommands(self, 417 cmdlist, 418 machine=None, 419 username=None, 420 command_terminator=None): 421 cmd = ' ;\n'.join(cmdlist) 422 return self.RunCommand(cmd, 423 machine=machine, 424 username=username, 425 command_terminator=command_terminator) 426 427 def CopyFiles(self, 428 src, 429 dest, 430 src_machine=None, 431 dest_machine=None, 432 src_user=None, 433 dest_user=None, 434 recursive=True, 435 command_terminator=None, 436 chromeos_root=None, 437 src_cros=False, 438 dest_cros=False, 439 print_to_console=True): 440 src = os.path.expanduser(src) 441 dest = os.path.expanduser(dest) 442 443 if recursive: 444 src = src + '/' 445 dest = dest + '/' 446 447 if src_cros == True or dest_cros == True: 448 if self.logger: 449 self.logger.LogFatalIf(src_cros == dest_cros, 450 'Only one of src_cros and desc_cros can ' 451 'be True.') 452 self.logger.LogFatalIf(not chromeos_root, 'chromeos_root not given!') 453 elif src_cros == dest_cros or not chromeos_root: 454 sys.exit(1) 455 if src_cros == True: 456 cros_machine = src_machine 457 else: 458 cros_machine = dest_machine 459 460 command = self.RemoteAccessInitCommand(chromeos_root, cros_machine) 461 ssh_command = ( 462 'ssh -p ${FLAGS_ssh_port}' + ' -o StrictHostKeyChecking=no' + 463 ' -o UserKnownHostsFile=$(mktemp)' + ' -i $TMP_PRIVATE_KEY') 464 rsync_prefix = "\nrsync -r -e \"%s\" " % ssh_command 465 if dest_cros == True: 466 command += rsync_prefix + '%s root@%s:%s' % (src, dest_machine, dest) 467 return self.RunCommand(command, 468 machine=src_machine, 469 username=src_user, 470 command_terminator=command_terminator, 471 print_to_console=print_to_console) 472 else: 473 command += rsync_prefix + 'root@%s:%s %s' % (src_machine, src, dest) 474 return self.RunCommand(command, 475 machine=dest_machine, 476 username=dest_user, 477 command_terminator=command_terminator, 478 print_to_console=print_to_console) 479 480 if dest_machine == src_machine: 481 command = 'rsync -a %s %s' % (src, dest) 482 else: 483 if src_machine is None: 484 src_machine = os.uname()[1] 485 src_user = getpass.getuser() 486 command = 'rsync -a %s@%s:%s %s' % (src_user, src_machine, src, dest) 487 return self.RunCommand(command, 488 machine=dest_machine, 489 username=dest_user, 490 command_terminator=command_terminator, 491 print_to_console=print_to_console) 492 493 def RunCommand2(self, 494 cmd, 495 cwd=None, 496 line_consumer=None, 497 timeout=None, 498 shell=True, 499 join_stderr=True, 500 env=None, 501 except_handler=lambda p, e: None): 502 """Run the command with an extra feature line_consumer. 503 504 This version allow developers to provide a line_consumer which will be 505 fed execution output lines. 506 507 A line_consumer is a callback, which is given a chance to run for each 508 line the execution outputs (either to stdout or stderr). The 509 line_consumer must accept one and exactly one dict argument, the dict 510 argument has these items - 511 'line' - The line output by the binary. Notice, this string includes 512 the trailing '\n'. 513 'output' - Whether this is a stdout or stderr output, values are either 514 'stdout' or 'stderr'. When join_stderr is True, this value 515 will always be 'output'. 516 'pobject' - The object used to control execution, for example, call 517 pobject.kill(). 518 519 Note: As this is written, the stdin for the process executed is 520 not associated with the stdin of the caller of this routine. 521 522 Args: 523 cmd: Command in a single string. 524 cwd: Working directory for execution. 525 line_consumer: A function that will ba called by this function. See above 526 for details. 527 timeout: terminate command after this timeout. 528 shell: Whether to use a shell for execution. 529 join_stderr: Whether join stderr to stdout stream. 530 env: Execution environment. 531 except_handler: Callback for when exception is thrown during command 532 execution. Passed process object and exception. 533 534 Returns: 535 Execution return code. 536 537 Raises: 538 child_exception: if fails to start the command process (missing 539 permission, no such file, etc) 540 """ 541 542 class StreamHandler(object): 543 """Internal utility class.""" 544 545 def __init__(self, pobject, fd, name, line_consumer): 546 self._pobject = pobject 547 self._fd = fd 548 self._name = name 549 self._buf = '' 550 self._line_consumer = line_consumer 551 552 def read_and_notify_line(self): 553 t = os.read(fd, 1024) 554 self._buf = self._buf + t 555 self.notify_line() 556 557 def notify_line(self): 558 p = self._buf.find('\n') 559 while p >= 0: 560 self._line_consumer(line=self._buf[:p + 1], 561 output=self._name, 562 pobject=self._pobject) 563 if p < len(self._buf) - 1: 564 self._buf = self._buf[p + 1:] 565 p = self._buf.find('\n') 566 else: 567 self._buf = '' 568 p = -1 569 break 570 571 def notify_eos(self): 572 # Notify end of stream. The last line may not end with a '\n'. 573 if self._buf != '': 574 self._line_consumer(line=self._buf, 575 output=self._name, 576 pobject=self._pobject) 577 self._buf = '' 578 579 if self.log_level == 'verbose': 580 self.logger.LogCmd(cmd) 581 elif self.logger: 582 self.logger.LogCmdToFileOnly(cmd) 583 584 # We use setsid so that the child will have a different session id 585 # and we can easily kill the process group. This is also important 586 # because the child will be disassociated from the parent terminal. 587 # In this way the child cannot mess the parent's terminal. 588 pobject = None 589 try: 590 pobject = subprocess.Popen( 591 cmd, 592 cwd=cwd, 593 bufsize=1024, 594 env=env, 595 shell=shell, 596 universal_newlines=True, 597 stdout=subprocess.PIPE, 598 stderr=subprocess.STDOUT if join_stderr else subprocess.PIPE, 599 preexec_fn=os.setsid) 600 601 # We provide a default line_consumer 602 if line_consumer is None: 603 line_consumer = lambda **d: None 604 start_time = time.time() 605 poll = select.poll() 606 outfd = pobject.stdout.fileno() 607 poll.register(outfd, select.POLLIN | select.POLLPRI) 608 handlermap = {outfd: 609 StreamHandler(pobject, outfd, 'stdout', line_consumer)} 610 if not join_stderr: 611 errfd = pobject.stderr.fileno() 612 poll.register(errfd, 613 select.POLLIN | select.POLLPRI) 614 handlermap[errfd] = StreamHandler(pobject, 615 errfd, 616 'stderr', 617 line_consumer) 618 while len(handlermap): 619 readables = poll.poll(300) 620 for (fd, evt) in readables: 621 handler = handlermap[fd] 622 if evt & (select.POLLPRI | select.POLLIN): 623 handler.read_and_notify_line() 624 elif evt & (select.POLLHUP | select.POLLERR | select.POLLNVAL): 625 handler.notify_eos() 626 poll.unregister(fd) 627 del handlermap[fd] 628 629 if timeout is not None and (time.time() - start_time > timeout): 630 os.killpg(os.getpgid(pobject.pid), signal.SIGTERM) 631 632 return pobject.wait() 633 except BaseException as e: 634 except_handler(pobject, e) 635 raise 636 637 638class MockCommandExecuter(CommandExecuter): 639 """Mock class for class CommandExecuter.""" 640 641 def __init__(self, log_level, logger_to_set=None): 642 super(MockCommandExecuter, self).__init__(log_level, logger_to_set) 643 644 def RunCommandGeneric(self, 645 cmd, 646 return_output=False, 647 machine=None, 648 username=None, 649 command_terminator=None, 650 command_timeout=None, 651 terminated_timeout=10, 652 print_to_console=True, 653 except_handler=lambda p, e: None): 654 assert not command_timeout 655 cmd = str(cmd) 656 if machine is None: 657 machine = 'localhost' 658 if username is None: 659 username = 'current' 660 logger.GetLogger().LogCmd('(Mock) ' + cmd, machine, username, 661 print_to_console) 662 return (0, '', '') 663 664 def RunCommand(self, *args, **kwargs): 665 assert 'return_output' not in kwargs 666 kwargs['return_output'] = False 667 return self.RunCommandGeneric(*args, **kwargs)[0] 668 669 def RunCommandWOutput(self, *args, **kwargs): 670 assert 'return_output' not in kwargs 671 kwargs['return_output'] = True 672 return self.RunCommandGeneric(*args, **kwargs) 673 674 675class CommandTerminator(object): 676 """Object to request termination of a command in execution.""" 677 678 def __init__(self): 679 self.terminated = False 680 681 def Terminate(self): 682 self.terminated = True 683 684 def IsTerminated(self): 685 return self.terminated 686