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