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