blob: 4f16c8506655b3ba53b1dd762c5f788a20be9b4f [file] [log] [blame]
Keun Soo Yimb293fdb2016-09-21 16:03:44 -07001#!/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 Deng26e4dc12018-03-04 19:01:59 -080016"""Common Utilities."""
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070017
Sam Chiu81bdc652018-06-29 18:45:08 +080018from __future__ import print_function
cylan66713722018-10-06 01:38:26 +080019
20from distutils.spawn import find_executable
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070021import base64
22import binascii
cylan66713722018-10-06 01:38:26 +080023import collections
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070024import errno
Fang Deng69498c32017-03-02 14:29:30 -080025import getpass
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070026import logging
27import os
28import shutil
29import struct
cylan66713722018-10-06 01:38:26 +080030import socket
Fang Deng69498c32017-03-02 14:29:30 -080031import subprocess
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070032import sys
33import tarfile
34import tempfile
35import time
Sam Chiu7a477f52018-10-22 11:20:36 +080036import Tkinter
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070037import uuid
chojoycecd004bc2018-09-13 10:39:00 +080038import zipfile
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070039
chojoycecd004bc2018-09-13 10:39:00 +080040from acloud import errors as root_errors
herbertxue34776bb2018-07-03 21:57:48 +080041from acloud.internal import constants
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070042from acloud.public import errors
43
44logger = logging.getLogger(__name__)
45
Fang Deng69498c32017-03-02 14:29:30 -080046SSH_KEYGEN_CMD = ["ssh-keygen", "-t", "rsa", "-b", "4096"]
cylan4f73c1f2018-07-19 16:40:31 +080047SSH_KEYGEN_PUB_CMD = ["ssh-keygen", "-y"]
Kevin Chengd25feee2018-05-24 10:15:20 -070048DEFAULT_RETRY_BACKOFF_FACTOR = 1
49DEFAULT_SLEEP_MULTIPLIER = 0
Fang Deng69498c32017-03-02 14:29:30 -080050
cylan66713722018-10-06 01:38:26 +080051_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")
cylan66713722018-10-06 01:38:26 +080056_ADB_CONNECT_ARGS = "connect 127.0.0.1:%(adb_port)d"
57# Store the ports that vnc/adb are forwarded to, both are integers.
58ForwardedPorts = collections.namedtuple("ForwardedPorts", [constants.VNC_PORT,
59 constants.ADB_PORT])
Kevin Chengae7a49d2018-10-18 14:11:22 -070060_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 Chiu7a477f52018-10-22 11:20:36 +080064_SSVNC_ENV_VARS = {"SSVNC_NO_ENC_WARN": "1", "SSVNC_SCALE": "auto"}
Kevin Chengae7a49d2018-10-18 14:11:22 -070065
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 Deng69498c32017-03-02 14:29:30 -080071
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070072class TempDir(object):
Fang Deng26e4dc12018-03-04 19:01:59 -080073 """A context manager that ceates a temporary directory.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070074
Fang Deng26e4dc12018-03-04 19:01:59 -080075 Attributes:
Sam Chiu81bdc652018-06-29 18:45:08 +080076 path: The path of the temporary directory.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070077 """
78
Fang Deng26e4dc12018-03-04 19:01:59 -080079 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 Yimb293fdb2016-09-21 16:03:44 -070083
84 def __enter__(self):
Fang Deng26e4dc12018-03-04 19:01:59 -080085 """Enter."""
86 return self.path
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070087
Fang Deng26e4dc12018-03-04 19:01:59 -080088 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 Yimb293fdb2016-09-21 16:03:44 -0700102 try:
Fang Deng26e4dc12018-03-04 19:01:59 -0800103 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 Yimb293fdb2016-09-21 16:03:44 -0700113 if exc_type:
Fang Deng26e4dc12018-03-04 19:01:59 -0800114 logger.error(
cylan0d77ae12018-05-18 08:36:48 +0000115 "Encountered error while deleting %s: %s",
116 self.path,
117 str(e),
118 exc_info=True)
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700119 else:
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700120 raise
121
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700122
cylan0d77ae12018-05-18 08:36:48 +0000123def RetryOnException(retry_checker,
124 max_retries,
125 sleep_multiplier=0,
Fang Dengf24be082018-02-10 10:09:55 -0800126 retry_backoff_factor=1):
cylan0d77ae12018-05-18 08:36:48 +0000127 """Decorater which retries the function call if |retry_checker| returns true.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700128
cylan0d77ae12018-05-18 08:36:48 +0000129 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 Yimb293fdb2016-09-21 16:03:44 -0700141
cylan0d77ae12018-05-18 08:36:48 +0000142 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 Yimb293fdb2016-09-21 16:03:44 -0700154
155
cylan4f73c1f2018-07-19 16:40:31 +0800156def Retry(retry_checker, max_retries, functor, sleep_multiplier,
157 retry_backoff_factor, *args, **kwargs):
cylan0d77ae12018-05-18 08:36:48 +0000158 """Conditionally retry a function.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700159
cylan0d77ae12018-05-18 08:36:48 +0000160 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 Yimb293fdb2016-09-21 16:03:44 -0700175
cylan0d77ae12018-05-18 08:36:48 +0000176 Returns:
177 The return value of the functor.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700178
cylan0d77ae12018-05-18 08:36:48 +0000179 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 Chengd25feee2018-05-24 10:15:20 -0700195 time.sleep(sleep)
cylan0d77ae12018-05-18 08:36:48 +0000196 else:
197 raise
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700198
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700199
Fang Dengf24be082018-02-10 10:09:55 -0800200def RetryExceptionType(exception_types, max_retries, functor, *args, **kwargs):
cylan0d77ae12018-05-18 08:36:48 +0000201 """Retry exception if it is one of the given types.
Fang Dengf24be082018-02-10 10:09:55 -0800202
cylan0d77ae12018-05-18 08:36:48 +0000203 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 Dengf24be082018-02-10 10:09:55 -0800210
cylan0d77ae12018-05-18 08:36:48 +0000211 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 Yimb293fdb2016-09-21 16:03:44 -0700216
217
218def 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
251def GenerateUniqueName(prefix=None, suffix=None):
Sam Chiu81bdc652018-06-29 18:45:08 +0800252 """Generate a random unique name using uuid4.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700253
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
269def 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 Deng69498c32017-03-02 14:29:30 -0800283def CreateSshKeyPairIfNotExist(private_key_path, public_key_path):
284 """Create the ssh key pair if they don't exist.
285
cylan4f73c1f2018-07-19 16:40:31 +0800286 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 Deng69498c32017-03-02 14:29:30 -0800292
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
cylan4f73c1f2018-07-19 16:40:31 +0800298
Fang Deng69498c32017-03-02 14:29:30 -0800299 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)
cylan4f73c1f2018-07-19 16:40:31 +0800304 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:
cylan0d77ae12018-05-18 08:36:48 +0000307 logger.debug(
cylan4f73c1f2018-07-19 16:40:31 +0800308 "The ssh private key (%s) and public key (%s) already exist,"
cylan0d77ae12018-05-18 08:36:48 +0000309 "will not automatically create the key pairs.", private_key_path,
310 public_key_path)
Fang Deng69498c32017-03-02 14:29:30 -0800311 return
cylan4f73c1f2018-07-19 16:40:31 +0800312 key_folder = os.path.dirname(private_key_path)
313 if not os.path.exists(key_folder):
314 os.makedirs(key_folder)
Fang Deng69498c32017-03-02 14:29:30 -0800315 try:
cylan4f73c1f2018-07-19 16:40:31 +0800316 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 Deng69498c32017-03-02 14:29:30 -0800334 except subprocess.CalledProcessError as e:
cylan0d77ae12018-05-18 08:36:48 +0000335 raise errors.DriverError("Failed to create ssh key pair: %s" % str(e))
Fang Deng69498c32017-03-02 14:29:30 -0800336 except OSError as e:
337 raise errors.DriverError(
cylan0d77ae12018-05-18 08:36:48 +0000338 "Failed to create ssh key pair, please make sure "
339 "'ssh-keygen' is installed: %s" % str(e))
Fang Deng69498c32017-03-02 14:29:30 -0800340
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(
cylan0d77ae12018-05-18 08:36:48 +0000350 "Failed to rename %s to %s: %s" % (default_pub_key_path,
351 public_key_path, str(e)))
Fang Deng69498c32017-03-02 14:29:30 -0800352
353 logger.info("Created ssh private key (%s) and public key (%s)",
354 private_key_path, public_key_path)
355
356
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700357def 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:
cylan0d77ae12018-05-18 08:36:48 +0000390 raise errors.DriverError(
391 "rsa key is invalid: %s, error: %s" % (rsa, str(e)))
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700392
chojoycecd004bc2018-09-13 10:39:00 +0800393def 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 Yimb293fdb2016-09-21 16:03:44 -0700415
Sam Chiu81bdc652018-06-29 18:45:08 +0800416# pylint: disable=old-style-class,no-init
417class 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
herbertxuedf01c422018-09-06 19:52:52 +0800430def PrintColorString(message, colors=TextColors.OKBLUE, **kwargs):
Sam Chiu81bdc652018-06-29 18:45:08 +0800431 """A helper function to print out colored text.
432
herbertxuedf01c422018-09-06 19:52:52 +0800433 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 Chiu81bdc652018-06-29 18:45:08 +0800443 Args:
444 message: String, the message text.
445 colors: String, color code.
herbertxuedf01c422018-09-06 19:52:52 +0800446 **kwargs: dictionary of keyword based args to pass to func.
Sam Chiu81bdc652018-06-29 18:45:08 +0800447 """
herbertxuedf01c422018-09-06 19:52:52 +0800448 print(colors + message + TextColors.ENDC, **kwargs)
449 sys.stdout.flush()
Sam Chiu81bdc652018-06-29 18:45:08 +0800450
451
452def 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
herbertxuedf01c422018-09-06 19:52:52 +0800464
herbertxue34776bb2018-07-03 21:57:48 +0800465def 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 Chiu81bdc652018-06-29 18:45:08 +0800478
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700479class 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 Dengf24be082018-02-10 10:09:55 -0800505 max_retry: See utils.Retry.
506 sleep: See utils.Retry.
507 backoff_factor: See utils.Retry.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700508 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 Chiu81bdc652018-06-29 18:45:08 +0800524 """Check if an exception is retriable.
525
526 Args:
527 exception: An exception instance.
528 """
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700529 if isinstance(exception, self._other_retriable_errors):
530 return True
531
cylan0d77ae12018-05-18 08:36:48 +0000532 if (isinstance(exception, errors.HttpError)
533 and exception.code in self._retry_http_codes):
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700534 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 Dengf24be082018-02-10 10:09:55 -0800555 # so that Retry will retry this function with pending_requests.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700556 raise errors.HasRetriableRequestsError(
cylan0d77ae12018-05-18 08:36:48 +0000557 "Retriable errors: %s" %
558 [str(results[rid][1]) for rid in self._pending_requests])
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700559
560 def Execute(self):
561 """Executes the requests and retry if necessary.
562
563 Will populate self._final_results.
564 """
cylan0d77ae12018-05-18 08:36:48 +0000565
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700566 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 Dengf24be082018-02-10 10:09:55 -0800583 Retry(
cylan0d77ae12018-05-18 08:36:48 +0000584 _ShouldRetryHandler,
585 max_retries=self._max_retry,
Fang Dengf24be082018-02-10 10:09:55 -0800586 functor=self._ExecuteOnce,
587 sleep_multiplier=self._sleep,
588 retry_backoff_factor=self._backoff_factor)
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700589 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
cylan31fc5332018-09-17 22:12:08 +0800603
604
605class 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
cylan66713722018-10-06 01:38:26 +0800653
654
655def 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
668def _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
688def 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 Cheng835a4152018-10-11 10:46:57 -0700715 _ExecuteCommand(constants.SSH_BIN, ssh_tunnel_args.split())
cylan66713722018-10-06 01:38:26 +0800716 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 Cheng835a4152018-10-11 10:46:57 -0700721 _ExecuteCommand(constants.ADB_BIN, adb_connect_args.split())
cylan66713722018-10-06 01:38:26 +0800722 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 Chengeb85e862018-10-09 15:35:13 -0700728
729
730def 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 Chengae7a49d2018-10-18 14:11:22 -0700766
767
Sam Chiu7a477f52018-10-22 11:20:36 +0800768def 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
781def _LaunchVncClient(port=constants.DEFAULT_VNC_PORT, avd_width=None,
782 avd_height=None):
Kevin Chengae7a49d2018-10-18 14:11:22 -0700783 """Launch ssvnc.
784
785 Args:
786 port: Integer, port number.
Sam Chiu7a477f52018-10-22 11:20:36 +0800787 avd_width: String, the width of avd.
788 avd_height: String, the height of avd.
Kevin Chengae7a49d2018-10-18 14:11:22 -0700789 """
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 Chiu7a477f52018-10-22 11:20:36 +0800812 # 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 Chengae7a49d2018-10-18 14:11:22 -0700820 subprocess.Popen(ssvnc_args.split(), env=ssvnc_env)
821
822
823def 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 Chiu7a477f52018-10-22 11:20:36 +0800853
854
855def 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