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