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, 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( 246 'w', encoding='utf-8', delete=False, prefix=os.uname()[1], 247 suffix='.sh') as f: 248 f.write('#!/bin/bash\n') 249 f.write(contents) 250 f.flush() 251 return f.name 252 253 def CrosLearnBoard(self, chromeos_root, machine): 254 command = self.RemoteAccessInitCommand(chromeos_root, machine) 255 command += '\nlearn_board' 256 command += '\necho ${FLAGS_board}' 257 retval, output, _ = self.RunCommandWOutput(command) 258 if self.logger: 259 self.logger.LogFatalIf(retval, 'learn_board command failed') 260 elif retval: 261 sys.exit(1) 262 return output.split()[-1] 263 264 def CrosRunCommandGeneric(self, 265 cmd, 266 return_output=False, 267 machine=None, 268 command_terminator=None, 269 chromeos_root=None, 270 command_timeout=None, 271 terminated_timeout=10, 272 print_to_console=True): 273 """Run a command on a ChromeOS box. 274 275 Returns triplet (returncode, stdout, stderr). 276 """ 277 278 if self.log_level != 'verbose': 279 print_to_console = False 280 281 if self.logger: 282 self.logger.LogCmd(cmd, print_to_console=print_to_console) 283 self.logger.LogFatalIf(not machine, 'No machine provided!') 284 self.logger.LogFatalIf(not chromeos_root, 'chromeos_root not given!') 285 else: 286 if not chromeos_root or not machine: 287 sys.exit(1) 288 chromeos_root = os.path.expanduser(chromeos_root) 289 290 port = None 291 if ':' in machine: 292 machine, port = machine.split(':') 293 # Write all commands to a file. 294 command_file = self.WriteToTempShFile(cmd) 295 retval = self.CopyFiles( 296 command_file, 297 command_file, 298 dest_machine=machine, 299 dest_port=port, 300 command_terminator=command_terminator, 301 chromeos_root=chromeos_root, 302 dest_cros=True, 303 recursive=False, 304 print_to_console=print_to_console) 305 if retval: 306 if self.logger: 307 self.logger.LogError('Could not run remote command on machine.' 308 ' Is the machine up?') 309 return (retval, '', '') 310 311 command = self.RemoteAccessInitCommand(chromeos_root, machine, port) 312 command += '\nremote_sh bash %s' % command_file 313 command += '\nl_retval=$?; echo "$REMOTE_OUT"; exit $l_retval' 314 retval = self.RunCommandGeneric( 315 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( 376 'w', 377 encoding='utf-8', 378 delete=False, 379 dir=os.path.join(chromeos_root, '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 if ret: 398 return (ret, '', '') 399 400 # Run command_file inside the chroot, making sure that any "~" is expanded 401 # by the shell inside the chroot, not outside. 402 command = ("cd %s; cros_sdk %s -- bash -c '%s/%s'" % 403 (chromeos_root, cros_sdk_options, CHROMEOS_SCRIPTS_DIR, 404 os.path.basename(command_file))) 405 ret = self.RunCommandGeneric( 406 command, 407 return_output, 408 command_terminator=command_terminator, 409 command_timeout=command_timeout, 410 terminated_timeout=terminated_timeout, 411 print_to_console=print_to_console, 412 env=env) 413 os.remove(command_file) 414 return ret 415 416 def ChrootRunCommand(self, *args, **kwargs): 417 """Runs a command within the chroot. 418 419 Takes the same arguments as ChrootRunCommandGeneric except for 420 return_output. 421 Returns a single value returncode. 422 """ 423 # Make sure that args does not overwrite 'return_output' 424 assert len(args) <= 2 425 assert 'return_output' not in kwargs 426 kwargs['return_output'] = False 427 return self.ChrootRunCommandGeneric(*args, **kwargs)[0] 428 429 def ChrootRunCommandWOutput(self, *args, **kwargs): 430 """Runs a command within the chroot. 431 432 Takes the same arguments as ChrootRunCommandGeneric except for 433 return_output. 434 Returns a triplet (returncode, stdout, stderr). 435 """ 436 # Make sure that args does not overwrite 'return_output' 437 assert len(args) <= 2 438 assert 'return_output' not in kwargs 439 kwargs['return_output'] = True 440 return self.ChrootRunCommandGeneric(*args, **kwargs) 441 442 def RunCommands(self, 443 cmdlist, 444 machine=None, 445 username=None, 446 command_terminator=None): 447 cmd = ' ;\n'.join(cmdlist) 448 return self.RunCommand( 449 cmd, 450 machine=machine, 451 username=username, 452 command_terminator=command_terminator) 453 454 def CopyFiles(self, 455 src, 456 dest, 457 src_machine=None, 458 src_port=None, 459 dest_machine=None, 460 dest_port=None, 461 src_user=None, 462 dest_user=None, 463 recursive=True, 464 command_terminator=None, 465 chromeos_root=None, 466 src_cros=False, 467 dest_cros=False, 468 print_to_console=True): 469 src = os.path.expanduser(src) 470 dest = os.path.expanduser(dest) 471 472 if recursive: 473 src = src + '/' 474 dest = dest + '/' 475 476 if src_cros or dest_cros: 477 if self.logger: 478 self.logger.LogFatalIf( 479 src_cros == dest_cros, 'Only one of src_cros and desc_cros can ' 480 'be True.') 481 self.logger.LogFatalIf(not chromeos_root, 'chromeos_root not given!') 482 elif src_cros == dest_cros or not chromeos_root: 483 sys.exit(1) 484 if src_cros: 485 cros_machine = src_machine 486 cros_port = src_port 487 host_machine = dest_machine 488 host_user = dest_user 489 else: 490 cros_machine = dest_machine 491 cros_port = dest_port 492 host_machine = src_machine 493 host_user = src_user 494 495 command = self.RemoteAccessInitCommand(chromeos_root, cros_machine, 496 cros_port) 497 ssh_command = ('ssh -o StrictHostKeyChecking=no' + 498 ' -o UserKnownHostsFile=$(mktemp)' + 499 ' -i $TMP_PRIVATE_KEY') 500 if cros_port: 501 ssh_command += ' -p %s' % cros_port 502 rsync_prefix = '\nrsync -r -e "%s" ' % ssh_command 503 if dest_cros: 504 command += rsync_prefix + '%s root@%s:%s' % (src, cros_machine, dest) 505 else: 506 command += rsync_prefix + 'root@%s:%s %s' % (cros_machine, src, dest) 507 508 return self.RunCommand( 509 command, 510 machine=host_machine, 511 username=host_user, 512 command_terminator=command_terminator, 513 print_to_console=print_to_console) 514 515 if dest_machine == src_machine: 516 command = 'rsync -a %s %s' % (src, dest) 517 else: 518 if src_machine is None: 519 src_machine = os.uname()[1] 520 src_user = getpass.getuser() 521 command = 'rsync -a %s@%s:%s %s' % (src_user, src_machine, src, dest) 522 return self.RunCommand( 523 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( 597 line=self._buf[:p + 1], output=self._name, pobject=self._pobject) 598 if p < len(self._buf) - 1: 599 self._buf = self._buf[p + 1:] 600 p = self._buf.find('\n') 601 else: 602 self._buf = '' 603 p = -1 604 break 605 606 def notify_eos(self): 607 # Notify end of stream. The last line may not end with a '\n'. 608 if self._buf != '': 609 self._line_consumer( 610 line=self._buf, output=self._name, pobject=self._pobject) 611 self._buf = '' 612 613 if self.log_level == 'verbose': 614 self.logger.LogCmd(cmd) 615 elif self.logger: 616 self.logger.LogCmdToFileOnly(cmd) 617 618 # We use setsid so that the child will have a different session id 619 # and we can easily kill the process group. This is also important 620 # because the child will be disassociated from the parent terminal. 621 # In this way the child cannot mess the parent's terminal. 622 pobject = None 623 try: 624 # pylint: disable=bad-option-value, subprocess-popen-preexec-fn 625 pobject = subprocess.Popen( 626 cmd, 627 cwd=cwd, 628 bufsize=1024, 629 env=env, 630 shell=shell, 631 universal_newlines=True, 632 stdout=subprocess.PIPE, 633 stderr=subprocess.STDOUT if join_stderr else subprocess.PIPE, 634 preexec_fn=os.setsid) 635 636 # We provide a default line_consumer 637 if line_consumer is None: 638 line_consumer = lambda **d: None 639 start_time = time.time() 640 poll = select.poll() 641 outfd = pobject.stdout.fileno() 642 poll.register(outfd, select.POLLIN | select.POLLPRI) 643 handlermap = { 644 outfd: StreamHandler(pobject, outfd, 'stdout', line_consumer) 645 } 646 if not join_stderr: 647 errfd = pobject.stderr.fileno() 648 poll.register(errfd, select.POLLIN | select.POLLPRI) 649 handlermap[errfd] = StreamHandler(pobject, errfd, 'stderr', 650 line_consumer) 651 while handlermap: 652 readables = poll.poll(300) 653 for (fd, evt) in readables: 654 handler = handlermap[fd] 655 if evt & (select.POLLPRI | select.POLLIN): 656 handler.read_and_notify_line() 657 elif evt & (select.POLLHUP | select.POLLERR | select.POLLNVAL): 658 handler.notify_eos() 659 poll.unregister(fd) 660 del handlermap[fd] 661 662 if timeout is not None and (time.time() - start_time > timeout): 663 os.killpg(os.getpgid(pobject.pid), signal.SIGTERM) 664 665 return pobject.wait() 666 except BaseException as err: 667 except_handler(pobject, err) 668 raise 669 670 671class MockCommandExecuter(CommandExecuter): 672 """Mock class for class CommandExecuter.""" 673 674 def RunCommandGeneric(self, 675 cmd, 676 return_output=False, 677 machine=None, 678 username=None, 679 command_terminator=None, 680 command_timeout=None, 681 terminated_timeout=10, 682 print_to_console=True, 683 env=None, 684 except_handler=lambda p, e: None): 685 assert not command_timeout 686 cmd = str(cmd) 687 if machine is None: 688 machine = 'localhost' 689 if username is None: 690 username = 'current' 691 logger.GetLogger().LogCmd('(Mock) ' + cmd, machine, username, 692 print_to_console) 693 return (0, '', '') 694 695 def RunCommand(self, *args, **kwargs): 696 assert 'return_output' not in kwargs 697 kwargs['return_output'] = False 698 return self.RunCommandGeneric(*args, **kwargs)[0] 699 700 def RunCommandWOutput(self, *args, **kwargs): 701 assert 'return_output' not in kwargs 702 kwargs['return_output'] = True 703 return self.RunCommandGeneric(*args, **kwargs) 704 705 706class CommandTerminator(object): 707 """Object to request termination of a command in execution.""" 708 709 def __init__(self): 710 self.terminated = False 711 712 def Terminate(self): 713 self.terminated = True 714 715 def IsTerminated(self): 716 return self.terminated 717