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