1# Copyright 2016 - The Android Open Source Project 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14"""Common Utilities.""" 15# pylint: disable=too-many-lines 16from __future__ import print_function 17 18from distutils.spawn import find_executable 19import base64 20import binascii 21import collections 22import errno 23import getpass 24import grp 25import logging 26import os 27import platform 28import shutil 29import signal 30import struct 31import socket 32import subprocess 33import sys 34import tarfile 35import tempfile 36import time 37import uuid 38import zipfile 39 40from acloud import errors 41from acloud.internal import constants 42 43logger = logging.getLogger(__name__) 44 45SSH_KEYGEN_CMD = ["ssh-keygen", "-t", "rsa", "-b", "4096"] 46SSH_KEYGEN_PUB_CMD = ["ssh-keygen", "-y"] 47SSH_ARGS = ["-o", "UserKnownHostsFile=/dev/null", 48 "-o", "StrictHostKeyChecking=no"] 49SSH_CMD = ["ssh"] + SSH_ARGS 50SCP_CMD = ["scp"] + SSH_ARGS 51GET_BUILD_VAR_CMD = ["build/soong/soong_ui.bash", "--dumpvar-mode"] 52DEFAULT_RETRY_BACKOFF_FACTOR = 1 53DEFAULT_SLEEP_MULTIPLIER = 0 54 55_SSH_TUNNEL_ARGS = ("-i %(rsa_key_file)s -o UserKnownHostsFile=/dev/null " 56 "-o StrictHostKeyChecking=no " 57 "-L %(vnc_port)d:127.0.0.1:%(target_vnc_port)d " 58 "-L %(adb_port)d:127.0.0.1:%(target_adb_port)d " 59 "-N -f -l %(ssh_user)s %(ip_addr)s") 60_ADB_CONNECT_ARGS = "connect 127.0.0.1:%(adb_port)d" 61# Store the ports that vnc/adb are forwarded to, both are integers. 62ForwardedPorts = collections.namedtuple("ForwardedPorts", [constants.VNC_PORT, 63 constants.ADB_PORT]) 64AVD_PORT_DICT = { 65 constants.TYPE_GCE: ForwardedPorts(constants.GCE_VNC_PORT, 66 constants.GCE_ADB_PORT), 67 constants.TYPE_CF: ForwardedPorts(constants.CF_VNC_PORT, 68 constants.CF_ADB_PORT), 69 constants.TYPE_GF: ForwardedPorts(constants.GF_VNC_PORT, 70 constants.GF_ADB_PORT), 71 constants.TYPE_CHEEPS: ForwardedPorts(constants.CHEEPS_VNC_PORT, 72 constants.CHEEPS_ADB_PORT) 73} 74 75_VNC_BIN = "ssvnc" 76_CMD_KILL = ["pkill", "-9", "-f"] 77_CMD_PGREP = "pgrep" 78_CMD_SG = "sg " 79_CMD_START_VNC = "%(bin)s vnc://127.0.0.1:%(port)d" 80_CMD_INSTALL_SSVNC = "sudo apt-get --assume-yes install ssvnc" 81_ENV_DISPLAY = "DISPLAY" 82_SSVNC_ENV_VARS = {"SSVNC_NO_ENC_WARN": "1", "SSVNC_SCALE": "auto", "VNCVIEWER_X11CURSOR": "1"} 83_DEFAULT_DISPLAY_SCALE = 1.0 84_DIST_DIR = "DIST_DIR" 85 86_CONFIRM_CONTINUE = ("In order to display the screen to the AVD, we'll need to " 87 "install a vnc client (ssvnc). \nWould you like acloud to " 88 "install it for you? (%s) \nPress 'y' to continue or " 89 "anything else to abort it[y/N]: ") % _CMD_INSTALL_SSVNC 90_EvaluatedResult = collections.namedtuple("EvaluatedResult", 91 ["is_result_ok", "result_message"]) 92# dict of supported system and their distributions. 93_SUPPORTED_SYSTEMS_AND_DISTS = {"Linux": ["Ubuntu", "Debian"]} 94_DEFAULT_TIMEOUT_ERR = "Function did not complete within %d secs." 95 96 97class TempDir(object): 98 """A context manager that ceates a temporary directory. 99 100 Attributes: 101 path: The path of the temporary directory. 102 """ 103 104 def __init__(self): 105 self.path = tempfile.mkdtemp() 106 os.chmod(self.path, 0o700) 107 logger.debug("Created temporary dir %s", self.path) 108 109 def __enter__(self): 110 """Enter.""" 111 return self.path 112 113 def __exit__(self, exc_type, exc_value, traceback): 114 """Exit. 115 116 Args: 117 exc_type: Exception type raised within the context manager. 118 None if no execption is raised. 119 exc_value: Exception instance raised within the context manager. 120 None if no execption is raised. 121 traceback: Traceback for exeception that is raised within 122 the context manager. 123 None if no execption is raised. 124 Raises: 125 EnvironmentError or OSError when failed to delete temp directory. 126 """ 127 try: 128 if self.path: 129 shutil.rmtree(self.path) 130 logger.debug("Deleted temporary dir %s", self.path) 131 except EnvironmentError as e: 132 # Ignore error if there is no exception raised 133 # within the with-clause and the EnvironementError is 134 # about problem that directory or file does not exist. 135 if not exc_type and e.errno != errno.ENOENT: 136 raise 137 except Exception as e: # pylint: disable=W0703 138 if exc_type: 139 logger.error( 140 "Encountered error while deleting %s: %s", 141 self.path, 142 str(e), 143 exc_info=True) 144 else: 145 raise 146 147 148def RetryOnException(retry_checker, 149 max_retries, 150 sleep_multiplier=0, 151 retry_backoff_factor=1): 152 """Decorater which retries the function call if |retry_checker| returns true. 153 154 Args: 155 retry_checker: A callback function which should take an exception instance 156 and return True if functor(*args, **kwargs) should be retried 157 when such exception is raised, and return False if it should 158 not be retried. 159 max_retries: Maximum number of retries allowed. 160 sleep_multiplier: Will sleep sleep_multiplier * attempt_count seconds if 161 retry_backoff_factor is 1. Will sleep 162 sleep_multiplier * ( 163 retry_backoff_factor ** (attempt_count - 1)) 164 if retry_backoff_factor != 1. 165 retry_backoff_factor: See explanation of sleep_multiplier. 166 167 Returns: 168 The function wrapper. 169 """ 170 171 def _Wrapper(func): 172 def _FunctionWrapper(*args, **kwargs): 173 return Retry(retry_checker, max_retries, func, sleep_multiplier, 174 retry_backoff_factor, *args, **kwargs) 175 176 return _FunctionWrapper 177 178 return _Wrapper 179 180 181def Retry(retry_checker, max_retries, functor, sleep_multiplier, 182 retry_backoff_factor, *args, **kwargs): 183 """Conditionally retry a function. 184 185 Args: 186 retry_checker: A callback function which should take an exception instance 187 and return True if functor(*args, **kwargs) should be retried 188 when such exception is raised, and return False if it should 189 not be retried. 190 max_retries: Maximum number of retries allowed. 191 functor: The function to call, will call functor(*args, **kwargs). 192 sleep_multiplier: Will sleep sleep_multiplier * attempt_count seconds if 193 retry_backoff_factor is 1. Will sleep 194 sleep_multiplier * ( 195 retry_backoff_factor ** (attempt_count - 1)) 196 if retry_backoff_factor != 1. 197 retry_backoff_factor: See explanation of sleep_multiplier. 198 *args: Arguments to pass to the functor. 199 **kwargs: Key-val based arguments to pass to the functor. 200 201 Returns: 202 The return value of the functor. 203 204 Raises: 205 Exception: The exception that functor(*args, **kwargs) throws. 206 """ 207 attempt_count = 0 208 while attempt_count <= max_retries: 209 try: 210 attempt_count += 1 211 return_value = functor(*args, **kwargs) 212 return return_value 213 except Exception as e: # pylint: disable=W0703 214 if retry_checker(e) and attempt_count <= max_retries: 215 if retry_backoff_factor != 1: 216 sleep = sleep_multiplier * (retry_backoff_factor** 217 (attempt_count - 1)) 218 else: 219 sleep = sleep_multiplier * attempt_count 220 time.sleep(sleep) 221 else: 222 raise 223 224 225def RetryExceptionType(exception_types, max_retries, functor, *args, **kwargs): 226 """Retry exception if it is one of the given types. 227 228 Args: 229 exception_types: A tuple of exception types, e.g. (ValueError, KeyError) 230 max_retries: Max number of retries allowed. 231 functor: The function to call. Will be retried if exception is raised and 232 the exception is one of the exception_types. 233 *args: Arguments to pass to Retry function. 234 **kwargs: Key-val based arguments to pass to Retry functions. 235 236 Returns: 237 The value returned by calling functor. 238 """ 239 return Retry(lambda e: isinstance(e, exception_types), max_retries, 240 functor, *args, **kwargs) 241 242 243def PollAndWait(func, expected_return, timeout_exception, timeout_secs, 244 sleep_interval_secs, *args, **kwargs): 245 """Call a function until the function returns expected value or times out. 246 247 Args: 248 func: Function to call. 249 expected_return: The expected return value. 250 timeout_exception: Exception to raise when it hits timeout. 251 timeout_secs: Timeout seconds. 252 If 0 or less than zero, the function will run once and 253 we will not wait on it. 254 sleep_interval_secs: Time to sleep between two attemps. 255 *args: list of args to pass to func. 256 **kwargs: dictionary of keyword based args to pass to func. 257 258 Raises: 259 timeout_exception: if the run of function times out. 260 """ 261 # TODO(fdeng): Currently this method does not kill 262 # |func|, if |func| takes longer than |timeout_secs|. 263 # We can use a more robust version from chromite. 264 start = time.time() 265 while True: 266 return_value = func(*args, **kwargs) 267 if return_value == expected_return: 268 return 269 elif time.time() - start > timeout_secs: 270 raise timeout_exception 271 else: 272 if sleep_interval_secs > 0: 273 time.sleep(sleep_interval_secs) 274 275 276def GenerateUniqueName(prefix=None, suffix=None): 277 """Generate a random unique name using uuid4. 278 279 Args: 280 prefix: String, desired prefix to prepend to the generated name. 281 suffix: String, desired suffix to append to the generated name. 282 283 Returns: 284 String, a random name. 285 """ 286 name = uuid.uuid4().hex 287 if prefix: 288 name = "-".join([prefix, name]) 289 if suffix: 290 name = "-".join([name, suffix]) 291 return name 292 293 294def MakeTarFile(src_dict, dest): 295 """Archive files in tar.gz format to a file named as |dest|. 296 297 Args: 298 src_dict: A dictionary that maps a path to be archived 299 to the corresponding name that appears in the archive. 300 dest: String, path to output file, e.g. /tmp/myfile.tar.gz 301 """ 302 logger.info("Compressing %s into %s.", src_dict.keys(), dest) 303 with tarfile.open(dest, "w:gz") as tar: 304 for src, arcname in src_dict.iteritems(): 305 tar.add(src, arcname=arcname) 306 307 308def ScpPullFile(src_file, dst_file, host_name, user_name=None, 309 rsa_key_file=None): 310 """Scp pull file from remote. 311 312 Args: 313 src_file: The source file path to be pulled. 314 dst_file: The destiation file path the file is pulled to. 315 host_name: The device host_name or ip to pull file from. 316 user_name: The user_name for scp session. 317 rsa_key_file: The rsa key file. 318 Raises: 319 errors.DeviceConnectionError if scp failed. 320 """ 321 scp_cmd_list = SCP_CMD[:] 322 if rsa_key_file: 323 scp_cmd_list.extend(["-i", rsa_key_file]) 324 else: 325 logger.warning( 326 "Rsa key file is not specified. " 327 "Will use default rsa key set in user environment") 328 if user_name: 329 scp_cmd_list.append("%s@%s:%s" % (user_name, host_name, src_file)) 330 else: 331 scp_cmd_list.append("%s:%s" % (host_name, src_file)) 332 scp_cmd_list.append(dst_file) 333 try: 334 subprocess.check_call(scp_cmd_list) 335 except subprocess.CalledProcessError as e: 336 raise errors.DeviceConnectionError( 337 "Failed to pull file %s from %s with '%s': %s" % ( 338 src_file, host_name, " ".join(scp_cmd_list), e)) 339 340 341def CreateSshKeyPairIfNotExist(private_key_path, public_key_path): 342 """Create the ssh key pair if they don't exist. 343 344 Case1. If the private key doesn't exist, we will create both the public key 345 and the private key. 346 Case2. If the private key exists but public key doesn't, we will create the 347 public key by using the private key. 348 Case3. If the public key exists but the private key doesn't, we will create 349 a new private key and overwrite the public key. 350 351 Args: 352 private_key_path: Path to the private key file. 353 e.g. ~/.ssh/acloud_rsa 354 public_key_path: Path to the public key file. 355 e.g. ~/.ssh/acloud_rsa.pub 356 357 Raises: 358 error.DriverError: If failed to create the key pair. 359 """ 360 public_key_path = os.path.expanduser(public_key_path) 361 private_key_path = os.path.expanduser(private_key_path) 362 public_key_exist = os.path.exists(public_key_path) 363 private_key_exist = os.path.exists(private_key_path) 364 if public_key_exist and private_key_exist: 365 logger.debug( 366 "The ssh private key (%s) and public key (%s) already exist," 367 "will not automatically create the key pairs.", private_key_path, 368 public_key_path) 369 return 370 key_folder = os.path.dirname(private_key_path) 371 if not os.path.exists(key_folder): 372 os.makedirs(key_folder) 373 try: 374 if private_key_exist: 375 cmd = SSH_KEYGEN_PUB_CMD + ["-f", private_key_path] 376 with open(public_key_path, 'w') as outfile: 377 stream_content = subprocess.check_output(cmd) 378 outfile.write( 379 stream_content.rstrip('\n') + " " + getpass.getuser()) 380 logger.info( 381 "The ssh public key (%s) do not exist, " 382 "automatically creating public key, calling: %s", 383 public_key_path, " ".join(cmd)) 384 else: 385 cmd = SSH_KEYGEN_CMD + [ 386 "-C", getpass.getuser(), "-f", private_key_path 387 ] 388 logger.info( 389 "Creating public key from private key (%s) via cmd: %s", 390 private_key_path, " ".join(cmd)) 391 subprocess.check_call(cmd, stdout=sys.stderr, stderr=sys.stdout) 392 except subprocess.CalledProcessError as e: 393 raise errors.DriverError("Failed to create ssh key pair: %s" % str(e)) 394 except OSError as e: 395 raise errors.DriverError( 396 "Failed to create ssh key pair, please make sure " 397 "'ssh-keygen' is installed: %s" % str(e)) 398 399 # By default ssh-keygen will create a public key file 400 # by append .pub to the private key file name. Rename it 401 # to what's requested by public_key_path. 402 default_pub_key_path = "%s.pub" % private_key_path 403 try: 404 if default_pub_key_path != public_key_path: 405 os.rename(default_pub_key_path, public_key_path) 406 except OSError as e: 407 raise errors.DriverError( 408 "Failed to rename %s to %s: %s" % (default_pub_key_path, 409 public_key_path, str(e))) 410 411 logger.info("Created ssh private key (%s) and public key (%s)", 412 private_key_path, public_key_path) 413 414 415def VerifyRsaPubKey(rsa): 416 """Verify the format of rsa public key. 417 418 Args: 419 rsa: content of rsa public key. It should follow the format of 420 ssh-rsa AAAAB3NzaC1yc2EA.... test@test.com 421 422 Raises: 423 DriverError if the format is not correct. 424 """ 425 if not rsa or not all(ord(c) < 128 for c in rsa): 426 raise errors.DriverError( 427 "rsa key is empty or contains non-ascii character: %s" % rsa) 428 429 elements = rsa.split() 430 if len(elements) != 3: 431 raise errors.DriverError("rsa key is invalid, wrong format: %s" % rsa) 432 433 key_type, data, _ = elements 434 try: 435 binary_data = base64.decodestring(data) 436 # number of bytes of int type 437 int_length = 4 438 # binary_data is like "7ssh-key..." in a binary format. 439 # The first 4 bytes should represent 7, which should be 440 # the length of the following string "ssh-key". 441 # And the next 7 bytes should be string "ssh-key". 442 # We will verify that the rsa conforms to this format. 443 # ">I" in the following line means "big-endian unsigned integer". 444 type_length = struct.unpack(">I", binary_data[:int_length])[0] 445 if binary_data[int_length:int_length + type_length] != key_type: 446 raise errors.DriverError("rsa key is invalid: %s" % rsa) 447 except (struct.error, binascii.Error) as e: 448 raise errors.DriverError( 449 "rsa key is invalid: %s, error: %s" % (rsa, str(e))) 450 451 452def Decompress(sourcefile, dest=None): 453 """Decompress .zip or .tar.gz. 454 455 Args: 456 sourcefile: A string, a source file path to decompress. 457 dest: A string, a folder path as decompress destination. 458 459 Raises: 460 errors.UnsupportedCompressionFileType: Not supported extension. 461 """ 462 logger.info("Start to decompress %s!", sourcefile) 463 dest_path = dest if dest else "." 464 if sourcefile.endswith(".tar.gz"): 465 with tarfile.open(sourcefile, "r:gz") as compressor: 466 compressor.extractall(dest_path) 467 elif sourcefile.endswith(".zip"): 468 with zipfile.ZipFile(sourcefile, 'r') as compressor: 469 compressor.extractall(dest_path) 470 else: 471 raise errors.UnsupportedCompressionFileType( 472 "Sorry, we could only support compression file type " 473 "for zip or tar.gz.") 474 475 476# pylint: disable=old-style-class,no-init 477class TextColors: 478 """A class that defines common color ANSI code.""" 479 480 HEADER = "\033[95m" 481 OKBLUE = "\033[94m" 482 OKGREEN = "\033[92m" 483 WARNING = "\033[33m" 484 FAIL = "\033[91m" 485 ENDC = "\033[0m" 486 BOLD = "\033[1m" 487 UNDERLINE = "\033[4m" 488 489 490def PrintColorString(message, colors=TextColors.OKBLUE, **kwargs): 491 """A helper function to print out colored text. 492 493 Use print function "print(message, end="")" to show message in one line. 494 Example code: 495 DisplayMessages("Creating GCE instance...", end="") 496 # Job execute 20s 497 DisplayMessages("Done! (20s)") 498 Display: 499 Creating GCE instance... 500 # After job finished, messages update as following: 501 Creating GCE instance...Done! (20s) 502 503 Args: 504 message: String, the message text. 505 colors: String, color code. 506 **kwargs: dictionary of keyword based args to pass to func. 507 """ 508 print(colors + message + TextColors.ENDC, **kwargs) 509 sys.stdout.flush() 510 511 512def InteractWithQuestion(question, colors=TextColors.WARNING): 513 """A helper function to define the common way to run interactive cmd. 514 515 Args: 516 question: String, the question to ask user. 517 colors: String, color code. 518 519 Returns: 520 String, input from user. 521 """ 522 return str(raw_input(colors + question + TextColors.ENDC).strip()) 523 524 525def GetUserAnswerYes(question): 526 """Ask user about acloud setup question. 527 528 Args: 529 question: String of question for user. Enter is equivalent to pressing 530 n. We should hint user with upper case N surrounded in square 531 brackets. 532 Ex: "Are you sure to change bucket name[y/N]:" 533 534 Returns: 535 Boolean, True if answer is "Yes", False otherwise. 536 """ 537 answer = InteractWithQuestion(question) 538 return answer.lower() in constants.USER_ANSWER_YES 539 540 541class BatchHttpRequestExecutor(object): 542 """A helper class that executes requests in batch with retry. 543 544 This executor executes http requests in a batch and retry 545 those that have failed. It iteratively updates the dictionary 546 self._final_results with latest results, which can be retrieved 547 via GetResults. 548 """ 549 550 def __init__(self, 551 execute_once_functor, 552 requests, 553 retry_http_codes=None, 554 max_retry=None, 555 sleep=None, 556 backoff_factor=None, 557 other_retriable_errors=None): 558 """Initializes the executor. 559 560 Args: 561 execute_once_functor: A function that execute requests in batch once. 562 It should return a dictionary like 563 {request_id: (response, exception)} 564 requests: A dictionary where key is request id picked by caller, 565 and value is a apiclient.http.HttpRequest. 566 retry_http_codes: A list of http codes to retry. 567 max_retry: See utils.Retry. 568 sleep: See utils.Retry. 569 backoff_factor: See utils.Retry. 570 other_retriable_errors: A tuple of error types that should be retried 571 other than errors.HttpError. 572 """ 573 self._execute_once_functor = execute_once_functor 574 self._requests = requests 575 # A dictionary that maps request id to pending request. 576 self._pending_requests = {} 577 # A dictionary that maps request id to a tuple (response, exception). 578 self._final_results = {} 579 self._retry_http_codes = retry_http_codes 580 self._max_retry = max_retry 581 self._sleep = sleep 582 self._backoff_factor = backoff_factor 583 self._other_retriable_errors = other_retriable_errors 584 585 def _ShoudRetry(self, exception): 586 """Check if an exception is retriable. 587 588 Args: 589 exception: An exception instance. 590 """ 591 if isinstance(exception, self._other_retriable_errors): 592 return True 593 594 if (isinstance(exception, errors.HttpError) 595 and exception.code in self._retry_http_codes): 596 return True 597 return False 598 599 def _ExecuteOnce(self): 600 """Executes pending requests and update it with failed, retriable ones. 601 602 Raises: 603 HasRetriableRequestsError: if some requests fail and are retriable. 604 """ 605 results = self._execute_once_functor(self._pending_requests) 606 # Update final_results with latest results. 607 self._final_results.update(results) 608 # Clear pending_requests 609 self._pending_requests.clear() 610 for request_id, result in results.iteritems(): 611 exception = result[1] 612 if exception is not None and self._ShoudRetry(exception): 613 # If this is a retriable exception, put it in pending_requests 614 self._pending_requests[request_id] = self._requests[request_id] 615 if self._pending_requests: 616 # If there is still retriable requests pending, raise an error 617 # so that Retry will retry this function with pending_requests. 618 raise errors.HasRetriableRequestsError( 619 "Retriable errors: %s" % 620 [str(results[rid][1]) for rid in self._pending_requests]) 621 622 def Execute(self): 623 """Executes the requests and retry if necessary. 624 625 Will populate self._final_results. 626 """ 627 628 def _ShouldRetryHandler(exc): 629 """Check if |exc| is a retriable exception. 630 631 Args: 632 exc: An exception. 633 634 Returns: 635 True if exception is of type HasRetriableRequestsError; False otherwise. 636 """ 637 should_retry = isinstance(exc, errors.HasRetriableRequestsError) 638 if should_retry: 639 logger.info("Will retry failed requests.", exc_info=True) 640 logger.info("%s", exc) 641 return should_retry 642 643 try: 644 self._pending_requests = self._requests.copy() 645 Retry( 646 _ShouldRetryHandler, 647 max_retries=self._max_retry, 648 functor=self._ExecuteOnce, 649 sleep_multiplier=self._sleep, 650 retry_backoff_factor=self._backoff_factor) 651 except errors.HasRetriableRequestsError: 652 logger.debug("Some requests did not succeed after retry.") 653 654 def GetResults(self): 655 """Returns final results. 656 657 Returns: 658 results, a dictionary in the following format 659 {request_id: (response, exception)} 660 request_ids are those from requests; response 661 is the http response for the request or None on error; 662 exception is an instance of DriverError or None if no error. 663 """ 664 return self._final_results 665 666 667def DefaultEvaluator(result): 668 """Default Evaluator always return result is ok. 669 670 Args: 671 result:the return value of the target function. 672 673 Returns: 674 _EvaluatedResults namedtuple. 675 """ 676 return _EvaluatedResult(is_result_ok=True, result_message=result) 677 678 679def ReportEvaluator(report): 680 """Evalute the acloud operation by the report. 681 682 Args: 683 report: acloud.public.report() object. 684 685 Returns: 686 _EvaluatedResults namedtuple. 687 """ 688 if report is None or report.errors: 689 return _EvaluatedResult(is_result_ok=False, 690 result_message=report.errors) 691 692 return _EvaluatedResult(is_result_ok=True, result_message=None) 693 694 695def BootEvaluator(boot_dict): 696 """Evaluate if the device booted successfully. 697 698 Args: 699 boot_dict: Dict of instance_name:boot error. 700 701 Returns: 702 _EvaluatedResults namedtuple. 703 """ 704 if boot_dict: 705 return _EvaluatedResult(is_result_ok=False, result_message=boot_dict) 706 return _EvaluatedResult(is_result_ok=True, result_message=None) 707 708 709class TimeExecute(object): 710 """Count the function execute time.""" 711 712 def __init__(self, function_description=None, print_before_call=True, 713 print_status=True, result_evaluator=DefaultEvaluator, 714 display_waiting_dots=True): 715 """Initializes the class. 716 717 Args: 718 function_description: String that describes function (e.g."Creating 719 Instance...") 720 print_before_call: Boolean, print the function description before 721 calling the function, default True. 722 print_status: Boolean, print the status of the function after the 723 function has completed, default True ("OK" or "Fail"). 724 result_evaluator: Func object. Pass func to evaluate result. 725 Default evaluator always report result is ok and 726 failed result will be identified only in exception 727 case. 728 display_waiting_dots: Boolean, if true print the function_description 729 followed by waiting dot. 730 """ 731 self._function_description = function_description 732 self._print_before_call = print_before_call 733 self._print_status = print_status 734 self._result_evaluator = result_evaluator 735 self._display_waiting_dots = display_waiting_dots 736 737 def __call__(self, func): 738 def DecoratorFunction(*args, **kargs): 739 """Decorator function. 740 741 Args: 742 *args: Arguments to pass to the functor. 743 **kwargs: Key-val based arguments to pass to the functor. 744 745 Raises: 746 Exception: The exception that functor(*args, **kwargs) throws. 747 """ 748 timestart = time.time() 749 if self._print_before_call: 750 waiting_dots = "..." if self._display_waiting_dots else "" 751 PrintColorString("%s %s"% (self._function_description, 752 waiting_dots), end="") 753 try: 754 result = func(*args, **kargs) 755 result_time = time.time() - timestart 756 if not self._print_before_call: 757 PrintColorString("%s (%ds)" % (self._function_description, 758 result_time), 759 TextColors.OKGREEN) 760 if self._print_status: 761 evaluated_result = self._result_evaluator(result) 762 if evaluated_result.is_result_ok: 763 PrintColorString("OK! (%ds)" % (result_time), 764 TextColors.OKGREEN) 765 else: 766 PrintColorString("Fail! (%ds)" % (result_time), 767 TextColors.FAIL) 768 PrintColorString("Error: %s" % 769 evaluated_result.result_message, 770 TextColors.FAIL) 771 return result 772 except: 773 if self._print_status: 774 PrintColorString("Fail! (%ds)" % (time.time()-timestart), 775 TextColors.FAIL) 776 raise 777 return DecoratorFunction 778 779 780def PickFreePort(): 781 """Helper to pick a free port. 782 783 Returns: 784 Integer, a free port number. 785 """ 786 tcp_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) 787 tcp_socket.bind(("", 0)) 788 port = tcp_socket.getsockname()[1] 789 tcp_socket.close() 790 return port 791 792 793def _ExecuteCommand(cmd, args): 794 """Execute command. 795 796 Args: 797 cmd: Strings of execute binary name. 798 args: List of args to pass in with cmd. 799 800 Raises: 801 errors.NoExecuteBin: Can't find the execute bin file. 802 """ 803 bin_path = find_executable(cmd) 804 if not bin_path: 805 raise errors.NoExecuteCmd("unable to locate %s" % cmd) 806 command = [bin_path] + args 807 logger.debug("Running '%s'", ' '.join(command)) 808 with open(os.devnull, "w") as dev_null: 809 subprocess.check_call(command, stderr=dev_null, stdout=dev_null) 810 811 812# pylint: disable=too-many-locals 813def AutoConnect(ip_addr, rsa_key_file, target_vnc_port, target_adb_port, ssh_user): 814 """Autoconnect to an AVD instance. 815 816 Args: 817 ip_addr: String, use to build the adb & vnc tunnel between local 818 and remote instance. 819 rsa_key_file: String, Private key file path to use when creating 820 the ssh tunnels. 821 target_vnc_port: Integer of target vnc port number. 822 target_adb_port: Integer of target adb port number. 823 ssh_user: String of user login into the instance. 824 825 Returns: 826 NamedTuple of (vnc_port, adb_port) SSHTUNNEL of the connect, both are 827 integers. 828 """ 829 local_free_vnc_port = PickFreePort() 830 local_free_adb_port = PickFreePort() 831 try: 832 ssh_tunnel_args = _SSH_TUNNEL_ARGS % { 833 "rsa_key_file": rsa_key_file, 834 "vnc_port": local_free_vnc_port, 835 "adb_port": local_free_adb_port, 836 "target_vnc_port": target_vnc_port, 837 "target_adb_port": target_adb_port, 838 "ssh_user": ssh_user, 839 "ip_addr": ip_addr} 840 _ExecuteCommand(constants.SSH_BIN, ssh_tunnel_args.split()) 841 except subprocess.CalledProcessError: 842 PrintColorString("Failed to create ssh tunnels, retry with '#acloud " 843 "reconnect'.", TextColors.FAIL) 844 return ForwardedPorts(vnc_port=None, adb_port=None) 845 846 try: 847 adb_connect_args = _ADB_CONNECT_ARGS % {"adb_port": local_free_adb_port} 848 _ExecuteCommand(constants.ADB_BIN, adb_connect_args.split()) 849 except subprocess.CalledProcessError: 850 PrintColorString("Failed to adb connect, retry with " 851 "'#acloud reconnect'", TextColors.FAIL) 852 853 return ForwardedPorts(vnc_port=local_free_vnc_port, 854 adb_port=local_free_adb_port) 855 856 857def GetAnswerFromList(answer_list, enable_choose_all=False): 858 """Get answer from a list. 859 860 Args: 861 answer_list: list of the answers to choose from. 862 enable_choose_all: True to choose all items from answer list. 863 864 Return: 865 List holding the answer(s). 866 """ 867 print("[0] to exit.") 868 start_index = 1 869 max_choice = len(answer_list) 870 871 for num, item in enumerate(answer_list, start_index): 872 print("[%d] %s" % (num, item)) 873 if enable_choose_all: 874 max_choice += 1 875 print("[%d] for all." % max_choice) 876 877 choice = -1 878 879 while True: 880 try: 881 choice = raw_input("Enter your choice[0-%d]: " % max_choice) 882 choice = int(choice) 883 except ValueError: 884 print("'%s' is not a valid integer.", choice) 885 continue 886 # Filter out choices 887 if choice == 0: 888 sys.exit(constants.EXIT_BY_USER) 889 if enable_choose_all and choice == max_choice: 890 return answer_list 891 if choice < 0 or choice > max_choice: 892 print("please choose between 0 and %d" % max_choice) 893 else: 894 return [answer_list[choice-start_index]] 895 896 897def LaunchVNCFromReport(report, avd_spec, no_prompts=False): 898 """Launch vnc client according to the instances report. 899 900 Args: 901 report: Report object, that stores and generates report. 902 avd_spec: AVDSpec object that tells us what we're going to create. 903 no_prompts: Boolean, True to skip all prompts. 904 """ 905 for device in report.data.get("devices", []): 906 if device.get(constants.VNC_PORT): 907 LaunchVncClient(device.get(constants.VNC_PORT), 908 avd_width=avd_spec.hw_property["x_res"], 909 avd_height=avd_spec.hw_property["y_res"], 910 no_prompts=no_prompts) 911 else: 912 PrintColorString("No VNC port specified, skipping VNC startup.", 913 TextColors.FAIL) 914 915 916def LaunchVncClient(port, avd_width=None, avd_height=None, no_prompts=False): 917 """Launch ssvnc. 918 919 Args: 920 port: Integer, port number. 921 avd_width: String, the width of avd. 922 avd_height: String, the height of avd. 923 no_prompts: Boolean, True to skip all prompts. 924 """ 925 try: 926 os.environ[_ENV_DISPLAY] 927 except KeyError: 928 PrintColorString("Remote terminal can't support VNC. " 929 "Skipping VNC startup.", TextColors.FAIL) 930 return 931 932 if not find_executable(_VNC_BIN): 933 if no_prompts or GetUserAnswerYes(_CONFIRM_CONTINUE): 934 try: 935 PrintColorString("Installing ssvnc vnc client... ", end="") 936 sys.stdout.flush() 937 subprocess.check_output(_CMD_INSTALL_SSVNC, shell=True) 938 PrintColorString("Done", TextColors.OKGREEN) 939 except subprocess.CalledProcessError as cpe: 940 PrintColorString("Failed to install ssvnc: %s" % 941 cpe.output, TextColors.FAIL) 942 return 943 else: 944 return 945 ssvnc_env = os.environ.copy() 946 ssvnc_env.update(_SSVNC_ENV_VARS) 947 # Override SSVNC_SCALE 948 if avd_width or avd_height: 949 scale_ratio = CalculateVNCScreenRatio(avd_width, avd_height) 950 ssvnc_env["SSVNC_SCALE"] = str(scale_ratio) 951 logger.debug("SSVNC_SCALE:%s", scale_ratio) 952 953 ssvnc_args = _CMD_START_VNC % {"bin": find_executable(_VNC_BIN), 954 "port": port} 955 subprocess.Popen(ssvnc_args.split(), env=ssvnc_env) 956 957 958def PrintDeviceSummary(report): 959 """Display summary of devices created. 960 961 -Display created device details from the report instance. 962 report example: 963 'data': [{'devices':[{'instance_name': 'ins-f6a397-none-53363', 964 'ip': u'35.234.10.162'}]}] 965 -Display error message from report.error. 966 967 Args: 968 report: A Report instance. 969 """ 970 PrintColorString("\n") 971 PrintColorString("Device(s) created:") 972 for device in report.data.get("devices", []): 973 adb_serial = "(None)" 974 adb_port = device.get("adb_port") 975 if adb_port: 976 adb_serial = constants.LOCALHOST_ADB_SERIAL % adb_port 977 instance_name = device.get("instance_name") 978 instance_ip = device.get("ip") 979 instance_details = "" if not instance_name else "(%s[%s])" % ( 980 instance_name, instance_ip) 981 PrintColorString(" - device serial: %s %s" % (adb_serial, 982 instance_details)) 983 984 # TODO(b/117245508): Help user to delete instance if it got created. 985 if report.errors: 986 error_msg = "\n".join(report.errors) 987 PrintColorString("Fail in:\n%s\n" % error_msg, TextColors.FAIL) 988 989 990def CalculateVNCScreenRatio(avd_width, avd_height): 991 """calculate the vnc screen scale ratio to fit into user's monitor. 992 993 Args: 994 avd_width: String, the width of avd. 995 avd_height: String, the height of avd. 996 Return: 997 Float, scale ratio for vnc client. 998 """ 999 try: 1000 import Tkinter 1001 # Some python interpreters may not be configured for Tk, just return default scale ratio. 1002 except ImportError: 1003 return _DEFAULT_DISPLAY_SCALE 1004 root = Tkinter.Tk() 1005 margin = 100 # leave some space on user's monitor. 1006 screen_height = root.winfo_screenheight() - margin 1007 screen_width = root.winfo_screenwidth() - margin 1008 1009 scale_h = _DEFAULT_DISPLAY_SCALE 1010 scale_w = _DEFAULT_DISPLAY_SCALE 1011 if float(screen_height) < float(avd_height): 1012 scale_h = round(float(screen_height) / float(avd_height), 1) 1013 1014 if float(screen_width) < float(avd_width): 1015 scale_w = round(float(screen_width) / float(avd_width), 1) 1016 1017 logger.debug("scale_h: %s (screen_h: %s/avd_h: %s)," 1018 " scale_w: %s (screen_w: %s/avd_w: %s)", 1019 scale_h, screen_height, avd_height, 1020 scale_w, screen_width, avd_width) 1021 1022 # Return the larger scale-down ratio. 1023 return scale_h if scale_h < scale_w else scale_w 1024 1025 1026def IsCommandRunning(command): 1027 """Check if command is running. 1028 1029 Args: 1030 command: String of command name. 1031 1032 Returns: 1033 Boolean, True if command is running. False otherwise. 1034 """ 1035 try: 1036 with open(os.devnull, "w") as dev_null: 1037 subprocess.check_call([_CMD_PGREP, "-f", command], 1038 stderr=dev_null, stdout=dev_null) 1039 return True 1040 except subprocess.CalledProcessError: 1041 return False 1042 1043 1044def AddUserGroupsToCmd(cmd, user_groups): 1045 """Add the user groups to the command if necessary. 1046 1047 As part of local host setup to enable local instance support, the user is 1048 added to certain groups. For those settings to take effect systemwide 1049 requires the user to log out and log back in. In the scenario where the 1050 user has run setup and hasn't logged out, we still want them to be able to 1051 launch a local instance so add the user to the groups as part of the 1052 command to ensure success. 1053 1054 The reason using here-doc instead of '&' is all operations need to be ran in 1055 ths same pid. Here's an example cmd: 1056 $ sg kvm << EOF 1057 sg libvirt 1058 sg cvdnetwork 1059 launch_cvd --cpus 2 --x_res 1280 --y_res 720 --dpi 160 --memory_mb 4096 1060 EOF 1061 1062 Args: 1063 cmd: String of the command to prepend the user groups to. 1064 user_groups: List of user groups name.(String) 1065 1066 Returns: 1067 String of the command with the user groups prepended to it if necessary, 1068 otherwise the same existing command. 1069 """ 1070 user_group_cmd = "" 1071 if not CheckUserInGroups(user_groups): 1072 logger.debug("Need to add user groups to the command") 1073 for idx, group in enumerate(user_groups): 1074 user_group_cmd += _CMD_SG + group 1075 if idx == 0: 1076 user_group_cmd += " <<EOF\n" 1077 else: 1078 user_group_cmd += "\n" 1079 cmd += "\nEOF" 1080 user_group_cmd += cmd 1081 logger.debug("user group cmd: %s", user_group_cmd) 1082 return user_group_cmd 1083 1084 1085def CheckUserInGroups(group_name_list): 1086 """Check if the current user is in the group. 1087 1088 Args: 1089 group_name_list: The list of group name. 1090 Returns: 1091 True if current user is in all the groups. 1092 """ 1093 logger.info("Checking if user is in following groups: %s", group_name_list) 1094 current_groups = [grp.getgrgid(g).gr_name for g in os.getgroups()] 1095 all_groups_present = True 1096 for group in group_name_list: 1097 if group not in current_groups: 1098 all_groups_present = False 1099 logger.info("missing group: %s", group) 1100 return all_groups_present 1101 1102 1103def IsSupportedPlatform(print_warning=False): 1104 """Check if user's os is the supported platform. 1105 1106 Args: 1107 print_warning: Boolean, print the unsupported warning 1108 if True. 1109 Returns: 1110 Boolean, True if user is using supported platform. 1111 """ 1112 system = platform.system() 1113 dist = platform.linux_distribution()[0] 1114 platform_supported = (system in _SUPPORTED_SYSTEMS_AND_DISTS and 1115 dist in _SUPPORTED_SYSTEMS_AND_DISTS[system]) 1116 1117 logger.info("supported system and dists: %s", 1118 _SUPPORTED_SYSTEMS_AND_DISTS) 1119 platform_supported_msg = ("%s[%s] %s supported platform" % 1120 (system, 1121 dist, 1122 "is a" if platform_supported else "is not a")) 1123 if print_warning and not platform_supported: 1124 PrintColorString(platform_supported_msg, TextColors.WARNING) 1125 else: 1126 logger.info(platform_supported_msg) 1127 1128 return platform_supported 1129 1130 1131def GetDistDir(): 1132 """Return the absolute path to the dist dir.""" 1133 android_build_top = os.environ.get(constants.ENV_ANDROID_BUILD_TOP) 1134 if not android_build_top: 1135 return None 1136 dist_cmd = GET_BUILD_VAR_CMD[:] 1137 dist_cmd.append(_DIST_DIR) 1138 try: 1139 dist_dir = subprocess.check_output(dist_cmd, cwd=android_build_top) 1140 except subprocess.CalledProcessError: 1141 return None 1142 return os.path.join(android_build_top, dist_dir.strip()) 1143 1144 1145def CleanupProcess(pattern): 1146 """Cleanup process with pattern. 1147 1148 Args: 1149 pattern: String, string of process pattern. 1150 """ 1151 if IsCommandRunning(pattern): 1152 command_kill = _CMD_KILL + [pattern] 1153 subprocess.check_call(command_kill) 1154 1155 1156def TimeoutException(timeout_secs, timeout_error=_DEFAULT_TIMEOUT_ERR): 1157 """Decorater which function timeout setup and raise custom exception. 1158 1159 Args: 1160 timeout_secs: Number of maximum seconds of waiting time. 1161 timeout_error: String to describe timeout exception. 1162 1163 Returns: 1164 The function wrapper. 1165 """ 1166 if timeout_error == _DEFAULT_TIMEOUT_ERR: 1167 timeout_error = timeout_error % timeout_secs 1168 1169 def _Wrapper(func): 1170 # pylint: disable=unused-argument 1171 def _HandleTimeout(signum, frame): 1172 raise errors.FunctionTimeoutError(timeout_error) 1173 1174 def _FunctionWrapper(*args, **kwargs): 1175 signal.signal(signal.SIGALRM, _HandleTimeout) 1176 signal.alarm(timeout_secs) 1177 try: 1178 result = func(*args, **kwargs) 1179 finally: 1180 signal.alarm(0) 1181 return result 1182 1183 return _FunctionWrapper 1184 1185 return _Wrapper 1186 1187 1188def GetBuildEnvironmentVariable(variable_name): 1189 """Get build environment variable. 1190 1191 Args: 1192 variable_name: String of variable name. 1193 1194 Returns: 1195 String, the value of the variable. 1196 1197 Raises: 1198 errors.GetAndroidBuildEnvVarError: No environment variable found. 1199 """ 1200 try: 1201 return os.environ[variable_name] 1202 except KeyError: 1203 raise errors.GetAndroidBuildEnvVarError( 1204 "Could not get environment var: %s\n" 1205 "Try to run 'source build/envsetup.sh && lunch <target>'" 1206 % variable_name 1207 ) 1208