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