1# Copyright 2016 Google Inc. 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 15import base64 16import concurrent.futures 17import datetime 18import errno 19import inspect 20import io 21import logging 22import os 23import platform 24import random 25import re 26import shlex 27import signal 28import string 29import subprocess 30import threading 31import time 32import traceback 33from typing import Literal, Tuple, overload 34 35import portpicker 36 37# File name length is limited to 255 chars on some OS, so we need to make sure 38# the file names we output fits within the limit. 39MAX_FILENAME_LEN = 255 40# Number of times to retry to get available port 41MAX_PORT_ALLOCATION_RETRY = 50 42 43ascii_letters_and_digits = string.ascii_letters + string.digits 44valid_filename_chars = f'-_.{ascii_letters_and_digits}' 45 46GMT_to_olson = { 47 'GMT-9': 'America/Anchorage', 48 'GMT-8': 'US/Pacific', 49 'GMT-7': 'US/Mountain', 50 'GMT-6': 'US/Central', 51 'GMT-5': 'US/Eastern', 52 'GMT-4': 'America/Barbados', 53 'GMT-3': 'America/Buenos_Aires', 54 'GMT-2': 'Atlantic/South_Georgia', 55 'GMT-1': 'Atlantic/Azores', 56 'GMT+0': 'Africa/Casablanca', 57 'GMT+1': 'Europe/Amsterdam', 58 'GMT+2': 'Europe/Athens', 59 'GMT+3': 'Europe/Moscow', 60 'GMT+4': 'Asia/Baku', 61 'GMT+5': 'Asia/Oral', 62 'GMT+6': 'Asia/Almaty', 63 'GMT+7': 'Asia/Bangkok', 64 'GMT+8': 'Asia/Hong_Kong', 65 'GMT+9': 'Asia/Tokyo', 66 'GMT+10': 'Pacific/Guam', 67 'GMT+11': 'Pacific/Noumea', 68 'GMT+12': 'Pacific/Fiji', 69 'GMT+13': 'Pacific/Tongatapu', 70 'GMT-11': 'Pacific/Midway', 71 'GMT-10': 'Pacific/Honolulu', 72} 73 74 75class Error(Exception): 76 """Raised when an error occurs in a util""" 77 78 79def abs_path(path): 80 """Resolve the '.' and '~' in a path to get the absolute path. 81 82 Args: 83 path: The path to expand. 84 85 Returns: 86 The absolute path of the input path. 87 """ 88 return os.path.abspath(os.path.expanduser(path)) 89 90 91def create_dir(path): 92 """Creates a directory if it does not exist already. 93 94 Args: 95 path: The path of the directory to create. 96 """ 97 full_path = abs_path(path) 98 if not os.path.exists(full_path): 99 try: 100 os.makedirs(full_path) 101 except OSError as e: 102 # ignore the error for dir already exist. 103 if e.errno != errno.EEXIST: 104 raise 105 106 107def create_alias(target_path, alias_path): 108 """Creates an alias at 'alias_path' pointing to the file 'target_path'. 109 110 On Unix, this is implemented via symlink. On Windows, this is done by 111 creating a Windows shortcut file. 112 113 Args: 114 target_path: Destination path that the alias should point to. 115 alias_path: Path at which to create the new alias. 116 """ 117 if platform.system() == 'Windows' and not alias_path.endswith('.lnk'): 118 alias_path += '.lnk' 119 if os.path.lexists(alias_path): 120 os.remove(alias_path) 121 if platform.system() == 'Windows': 122 from win32com import client 123 124 shell = client.Dispatch('WScript.Shell') 125 shortcut = shell.CreateShortCut(alias_path) 126 shortcut.Targetpath = target_path 127 shortcut.save() 128 else: 129 os.symlink(target_path, alias_path) 130 131 132def get_current_epoch_time(): 133 """Current epoch time in milliseconds. 134 135 Returns: 136 An integer representing the current epoch time in milliseconds. 137 """ 138 return int(round(time.time() * 1000)) 139 140 141def get_current_human_time(): 142 """Returns the current time in human readable format. 143 144 Returns: 145 The current time stamp in Month-Day-Year Hour:Min:Sec format. 146 """ 147 return time.strftime('%m-%d-%Y %H:%M:%S ') 148 149 150def epoch_to_human_time(epoch_time): 151 """Converts an epoch timestamp to human readable time. 152 153 This essentially converts an output of get_current_epoch_time to an output 154 of get_current_human_time 155 156 Args: 157 epoch_time: An integer representing an epoch timestamp in milliseconds. 158 159 Returns: 160 A time string representing the input time. 161 None if input param is invalid. 162 """ 163 if isinstance(epoch_time, int): 164 try: 165 d = datetime.datetime.fromtimestamp(epoch_time / 1000) 166 return d.strftime('%m-%d-%Y %H:%M:%S ') 167 except ValueError: 168 return None 169 170 171def get_timezone_olson_id(): 172 """Return the Olson ID of the local (non-DST) timezone. 173 174 Returns: 175 A string representing one of the Olson IDs of the local (non-DST) 176 timezone. 177 """ 178 tzoffset = int(time.timezone / 3600) 179 if tzoffset <= 0: 180 gmt = f'GMT+{-tzoffset}' 181 else: 182 gmt = f'GMT-{tzoffset}' 183 return GMT_to_olson[gmt] 184 185 186def find_files(paths, file_predicate): 187 """Locate files whose names and extensions match the given predicate in 188 the specified directories. 189 190 Args: 191 paths: A list of directory paths where to find the files. 192 file_predicate: A function that returns True if the file name and 193 extension are desired. 194 195 Returns: 196 A list of files that match the predicate. 197 """ 198 file_list = [] 199 for path in paths: 200 p = abs_path(path) 201 for dirPath, _, fileList in os.walk(p): 202 for fname in fileList: 203 name, ext = os.path.splitext(fname) 204 if file_predicate(name, ext): 205 file_list.append((dirPath, name, ext)) 206 return file_list 207 208 209def load_file_to_base64_str(f_path): 210 """Loads the content of a file into a base64 string. 211 212 Args: 213 f_path: full path to the file including the file name. 214 215 Returns: 216 A base64 string representing the content of the file in utf-8 encoding. 217 """ 218 path = abs_path(f_path) 219 with io.open(path, 'rb') as f: 220 f_bytes = f.read() 221 base64_str = base64.b64encode(f_bytes).decode('utf-8') 222 return base64_str 223 224 225def find_field(item_list, cond, comparator, target_field): 226 """Finds the value of a field in a dict object that satisfies certain 227 conditions. 228 229 Args: 230 item_list: A list of dict objects. 231 cond: A param that defines the condition. 232 comparator: A function that checks if an dict satisfies the condition. 233 target_field: Name of the field whose value to be returned if an item 234 satisfies the condition. 235 236 Returns: 237 Target value or None if no item satisfies the condition. 238 """ 239 for item in item_list: 240 if comparator(item, cond) and target_field in item: 241 return item[target_field] 242 return None 243 244 245def rand_ascii_str(length): 246 """Generates a random string of specified length, composed of ascii letters 247 and digits. 248 249 Args: 250 length: The number of characters in the string. 251 252 Returns: 253 The random string generated. 254 """ 255 letters = [random.choice(ascii_letters_and_digits) for _ in range(length)] 256 return ''.join(letters) 257 258 259# Thead/Process related functions. 260def _collect_process_tree(starting_pid): 261 """Collects PID list of the descendant processes from the given PID. 262 263 This function only available on Unix like system. 264 265 Args: 266 starting_pid: The PID to start recursively traverse. 267 268 Returns: 269 A list of pid of the descendant processes. 270 """ 271 ret = [] 272 stack = [starting_pid] 273 274 while stack: 275 pid = stack.pop() 276 if platform.system() == 'Darwin': 277 command = ['pgrep', '-P', str(pid)] 278 else: 279 command = [ 280 'ps', 281 '-o', 282 'pid', 283 '--ppid', 284 str(pid), 285 '--noheaders', 286 ] 287 try: 288 ps_results = subprocess.check_output(command).decode().strip() 289 except subprocess.CalledProcessError: 290 # Ignore if there is not child process. 291 continue 292 293 children_pid_list = [int(p.strip()) for p in ps_results.split('\n')] 294 stack.extend(children_pid_list) 295 ret.extend(children_pid_list) 296 297 return ret 298 299 300def _kill_process_tree(proc): 301 """Kills the subprocess and its descendants.""" 302 if os.name == 'nt': 303 # The taskkill command with "/T" option ends the specified process and any 304 # child processes started by it: 305 # https://docs.microsoft.com/en-us/windows-server/administration/windows-commands/taskkill 306 subprocess.check_output( 307 [ 308 'taskkill', 309 '/F', 310 '/T', 311 '/PID', 312 str(proc.pid), 313 ] 314 ) 315 return 316 317 failed = [] 318 for child_pid in _collect_process_tree(proc.pid): 319 try: 320 os.kill(child_pid, signal.SIGTERM) 321 except Exception: # pylint: disable=broad-except 322 failed.append(child_pid) 323 logging.exception('Failed to kill standing subprocess %d', child_pid) 324 325 try: 326 proc.kill() 327 except Exception: # pylint: disable=broad-except 328 failed.append(proc.pid) 329 logging.exception('Failed to kill standing subprocess %d', proc.pid) 330 331 if failed: 332 raise Error('Failed to kill standing subprocesses: %s' % failed) 333 334 335def concurrent_exec(func, param_list, max_workers=30, raise_on_exception=False): 336 """Executes a function with different parameters pseudo-concurrently. 337 338 This is basically a map function. Each element (should be an iterable) in 339 the param_list is unpacked and passed into the function. Due to Python's 340 GIL, there's no true concurrency. This is suited for IO-bound tasks. 341 342 Args: 343 func: The function that performs a task. 344 param_list: A list of iterables, each being a set of params to be 345 passed into the function. 346 max_workers: int, the number of workers to use for parallelizing the 347 tasks. By default, this is 30 workers. 348 raise_on_exception: bool, raises all of the task failures if any of the 349 tasks failed if `True`. By default, this is `False`. 350 351 Returns: 352 A list of return values from each function execution. If an execution 353 caused an exception, the exception object will be the corresponding 354 result. 355 356 Raises: 357 RuntimeError: If executing any of the tasks failed and 358 `raise_on_exception` is True. 359 """ 360 with concurrent.futures.ThreadPoolExecutor( 361 max_workers=max_workers 362 ) as executor: 363 # Start the load operations and mark each future with its params 364 future_to_params = {executor.submit(func, *p): p for p in param_list} 365 return_vals = [] 366 exceptions = [] 367 for future in concurrent.futures.as_completed(future_to_params): 368 params = future_to_params[future] 369 try: 370 return_vals.append(future.result()) 371 except Exception as exc: # pylint: disable=broad-except 372 logging.exception( 373 '%s generated an exception: %s', params, traceback.format_exc() 374 ) 375 return_vals.append(exc) 376 exceptions.append(exc) 377 if raise_on_exception and exceptions: 378 error_messages = [] 379 for exception in exceptions: 380 error_messages.append( 381 ''.join( 382 traceback.format_exception( 383 exception.__class__, exception, exception.__traceback__ 384 ) 385 ) 386 ) 387 raise RuntimeError('\n\n'.join(error_messages)) 388 return return_vals 389 390 391# Provide hint for pytype checker to avoid the Union[bytes, str] case. 392@overload 393def run_command( 394 cmd, 395 stdout=..., 396 stderr=..., 397 shell=..., 398 timeout=..., 399 cwd=..., 400 env=..., 401 universal_newlines: Literal[False] = ..., 402) -> Tuple[int, bytes, bytes]: 403 ... 404 405 406@overload 407def run_command( 408 cmd, 409 stdout=..., 410 stderr=..., 411 shell=..., 412 timeout=..., 413 cwd=..., 414 env=..., 415 universal_newlines: Literal[True] = ..., 416) -> Tuple[int, str, str]: 417 ... 418 419 420def run_command( 421 cmd, 422 stdout=None, 423 stderr=None, 424 shell=False, 425 timeout=None, 426 cwd=None, 427 env=None, 428 universal_newlines=False, 429): 430 """Runs a command in a subprocess. 431 432 This function is very similar to subprocess.check_output. The main 433 difference is that it returns the return code and std error output as well 434 as supporting a timeout parameter. 435 436 Args: 437 cmd: string or list of strings, the command to run. 438 See subprocess.Popen() documentation. 439 stdout: file handle, the file handle to write std out to. If None is 440 given, then subprocess.PIPE is used. See subprocess.Popen() 441 documentation. 442 stderr: file handle, the file handle to write std err to. If None is 443 given, then subprocess.PIPE is used. See subprocess.Popen() 444 documentation. 445 shell: bool, True to run this command through the system shell, 446 False to invoke it directly. See subprocess.Popen() docs. 447 timeout: float, the number of seconds to wait before timing out. 448 If not specified, no timeout takes effect. 449 cwd: string, the path to change the child's current directory to before 450 it is executed. Note that this directory is not considered when 451 searching the executable, so you can't specify the program's path 452 relative to cwd. 453 env: dict, a mapping that defines the environment variables for the 454 new process. Default behavior is inheriting the current process' 455 environment. 456 universal_newlines: bool, True to open file objects in text mode, False in 457 binary mode. 458 459 Returns: 460 A 3-tuple of the consisting of the return code, the std output, and the 461 std error. 462 463 Raises: 464 subprocess.TimeoutExpired: The command timed out. 465 """ 466 if stdout is None: 467 stdout = subprocess.PIPE 468 if stderr is None: 469 stderr = subprocess.PIPE 470 process = subprocess.Popen( 471 cmd, 472 stdout=stdout, 473 stderr=stderr, 474 shell=shell, 475 cwd=cwd, 476 env=env, 477 universal_newlines=universal_newlines, 478 ) 479 timer = None 480 timer_triggered = threading.Event() 481 if timeout and timeout > 0: 482 # The wait method on process will hang when used with PIPEs with large 483 # outputs, so use a timer thread instead. 484 485 def timeout_expired(): 486 timer_triggered.set() 487 process.terminate() 488 489 timer = threading.Timer(timeout, timeout_expired) 490 timer.start() 491 # If the command takes longer than the timeout, then the timer thread 492 # will kill the subprocess, which will make it terminate. 493 out, err = process.communicate() 494 if timer is not None: 495 timer.cancel() 496 if timer_triggered.is_set(): 497 raise subprocess.TimeoutExpired( 498 cmd=cmd, timeout=timeout, output=out, stderr=err 499 ) 500 logging.debug( 501 'cmd: %s, stdout: %s, stderr: %s, ret: %s', 502 cli_cmd_to_string(cmd), 503 out, 504 err, 505 process.returncode, 506 ) 507 return process.returncode, out, err 508 509 510def start_standing_subprocess( 511 cmd, 512 shell=False, 513 env=None, 514 stdout=subprocess.PIPE, 515 stderr=subprocess.PIPE, 516): 517 """Starts a long-running subprocess. 518 519 This is not a blocking call and the subprocess started by it should be 520 explicitly terminated with stop_standing_subprocess. 521 522 For short-running commands, you should use subprocess.check_call, which 523 blocks. 524 525 Args: 526 cmd: string, the command to start the subprocess with. 527 shell: bool, True to run this command through the system shell, 528 False to invoke it directly. See subprocess.Popen() docs. 529 env: dict, a custom environment to run the standing subprocess. If not 530 specified, inherits the current environment. See subprocess.Popen() 531 docs. 532 stdout: None, subprocess.PIPE, subprocess.DEVNULL, an existing file 533 descriptor, or an existing file object. See subprocess.Popen() docs. 534 stderr: None, subprocess.PIPE, subprocess.DEVNULL, an existing file 535 descriptor, or an existing file object. See subprocess.Popen() docs. 536 537 Returns: 538 The subprocess that was started. 539 """ 540 logging.debug('Starting standing subprocess with: %s', cmd) 541 proc = subprocess.Popen( 542 cmd, 543 stdin=subprocess.PIPE, 544 stdout=stdout, 545 stderr=stderr, 546 shell=shell, 547 env=env, 548 ) 549 # Leaving stdin open causes problems for input, e.g. breaking the 550 # code.inspect() shell (http://stackoverflow.com/a/25512460/1612937), so 551 # explicitly close it assuming it is not needed for standing subprocesses. 552 proc.stdin.close() 553 proc.stdin = None 554 logging.debug('Started standing subprocess %d', proc.pid) 555 return proc 556 557 558def stop_standing_subprocess(proc): 559 """Stops a subprocess started by start_standing_subprocess. 560 561 Before killing the process, we check if the process is running, if it has 562 terminated, Error is raised. 563 564 Catches and ignores the PermissionError which only happens on Macs. 565 566 Args: 567 proc: Subprocess to terminate. 568 569 Raises: 570 Error: if the subprocess could not be stopped. 571 """ 572 logging.debug('Stopping standing subprocess %d', proc.pid) 573 574 _kill_process_tree(proc) 575 576 # Call wait and close pipes on the original Python object so we don't get 577 # runtime warnings. 578 if proc.stdout: 579 proc.stdout.close() 580 if proc.stderr: 581 proc.stderr.close() 582 proc.wait() 583 logging.debug('Stopped standing subprocess %d', proc.pid) 584 585 586def wait_for_standing_subprocess(proc, timeout=None): 587 """Waits for a subprocess started by start_standing_subprocess to finish 588 or times out. 589 590 Propagates the exception raised by the subprocess.wait(.) function. 591 The subprocess.TimeoutExpired exception is raised if the process timed-out 592 rather than terminating. 593 594 If no exception is raised: the subprocess terminated on its own. No need 595 to call stop_standing_subprocess() to kill it. 596 597 If an exception is raised: the subprocess is still alive - it did not 598 terminate. Either call stop_standing_subprocess() to kill it, or call 599 wait_for_standing_subprocess() to keep waiting for it to terminate on its 600 own. 601 602 If the corresponding subprocess command generates a large amount of output 603 and this method is called with a timeout value, then the command can hang 604 indefinitely. See http://go/pylib/subprocess.html#subprocess.Popen.wait 605 606 This function does not support Python 2. 607 608 Args: 609 p: Subprocess to wait for. 610 timeout: An integer number of seconds to wait before timing out. 611 """ 612 proc.wait(timeout) 613 614 615def get_available_host_port(): 616 """Gets a host port number available for adb forward. 617 618 DEPRECATED: This method is unreliable. Pass `tcp:0` to adb forward instead. 619 620 Returns: 621 An integer representing a port number on the host available for adb 622 forward. 623 624 Raises: 625 Error: when no port is found after MAX_PORT_ALLOCATION_RETRY times. 626 """ 627 logging.warning( 628 'The method mobly.utils.get_available_host_port is deprecated because it ' 629 'is unreliable. Pass "tcp:0" to adb forward instead.' 630 ) 631 632 # Only import adb module if needed. 633 from mobly.controllers.android_device_lib import adb 634 635 port = portpicker.pick_unused_port() 636 if not adb.is_adb_available(): 637 return port 638 for _ in range(MAX_PORT_ALLOCATION_RETRY): 639 # Make sure adb is not using this port so we don't accidentally 640 # interrupt ongoing runs by trying to bind to the port. 641 if port not in adb.list_occupied_adb_ports(): 642 return port 643 port = portpicker.pick_unused_port() 644 raise Error( 645 'Failed to find available port after {} retries'.format( 646 MAX_PORT_ALLOCATION_RETRY 647 ) 648 ) 649 650 651def grep(regex, output): 652 """Similar to linux's `grep`, this returns the line in an output stream 653 that matches a given regex pattern. 654 655 It does not rely on the `grep` binary and is not sensitive to line endings, 656 so it can be used cross-platform. 657 658 Args: 659 regex: string, a regex that matches the expected pattern. 660 output: byte string, the raw output of the adb cmd. 661 662 Returns: 663 A list of strings, all of which are output lines that matches the 664 regex pattern. 665 """ 666 lines = output.decode('utf-8').strip().splitlines() 667 results = [] 668 for line in lines: 669 if re.search(regex, line): 670 results.append(line.strip()) 671 return results 672 673 674def cli_cmd_to_string(args): 675 """Converts a cmd arg list to string. 676 677 Args: 678 args: list of strings, the arguments of a command. 679 680 Returns: 681 String representation of the command. 682 """ 683 if isinstance(args, str): 684 # Return directly if it's already a string. 685 return args 686 return ' '.join([shlex.quote(arg) for arg in args]) 687 688 689def get_settable_properties(cls): 690 """Gets the settable properties of a class. 691 692 Only returns the explicitly defined properties with setters. 693 694 Args: 695 cls: A class in Python. 696 """ 697 results = [] 698 for attr, value in vars(cls).items(): 699 if isinstance(value, property) and value.fset is not None: 700 results.append(attr) 701 return results 702 703 704def find_subclasses_in_module(base_classes, module): 705 """Finds the subclasses of the given classes in the given module. 706 707 Args: 708 base_classes: list of classes, the base classes to look for the 709 subclasses of in the module. 710 module: module, the module to look for the subclasses in. 711 712 Returns: 713 A list of all of the subclasses found in the module. 714 """ 715 subclasses = [] 716 for _, module_member in module.__dict__.items(): 717 if inspect.isclass(module_member): 718 for base_class in base_classes: 719 if issubclass(module_member, base_class): 720 subclasses.append(module_member) 721 return subclasses 722 723 724def find_subclass_in_module(base_class, module): 725 """Finds the single subclass of the given base class in the given module. 726 727 Args: 728 base_class: class, the base class to look for a subclass of in the module. 729 module: module, the module to look for the single subclass in. 730 731 Returns: 732 The single subclass of the given base class. 733 734 Raises: 735 ValueError: If the number of subclasses found was not exactly one. 736 """ 737 subclasses = find_subclasses_in_module([base_class], module) 738 if len(subclasses) != 1: 739 raise ValueError( 740 'Expected 1 subclass of %s per module, found %s.' 741 % (base_class.__name__, [subclass.__name__ for subclass in subclasses]) 742 ) 743 return subclasses[0] 744