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