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