blob: 239716ca3baa9c5cd628d335655f12bd7c6e16e5 [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
36import uuid
chojoycecd004bc2018-09-13 10:39:00 +080037import zipfile
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070038
chojoycecd004bc2018-09-13 10:39:00 +080039from acloud import errors as root_errors
herbertxue34776bb2018-07-03 21:57:48 +080040from acloud.internal import constants
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070041from acloud.public import errors
42
43logger = logging.getLogger(__name__)
44
Fang Deng69498c32017-03-02 14:29:30 -080045SSH_KEYGEN_CMD = ["ssh-keygen", "-t", "rsa", "-b", "4096"]
cylan4f73c1f2018-07-19 16:40:31 +080046SSH_KEYGEN_PUB_CMD = ["ssh-keygen", "-y"]
Kevin Chengd25feee2018-05-24 10:15:20 -070047DEFAULT_RETRY_BACKOFF_FACTOR = 1
48DEFAULT_SLEEP_MULTIPLIER = 0
Fang Deng69498c32017-03-02 14:29:30 -080049
cylan66713722018-10-06 01:38:26 +080050_SSH_TUNNEL_ARGS = ("-i %(rsa_key_file)s -o UserKnownHostsFile=/dev/null "
51 "-o StrictHostKeyChecking=no "
52 "-L %(vnc_port)d:127.0.0.1:%(target_vnc_port)d "
53 "-L %(adb_port)d:127.0.0.1:%(target_adb_port)d "
54 "-N -f -l %(ssh_user)s %(ip_addr)s")
cylan66713722018-10-06 01:38:26 +080055_ADB_CONNECT_ARGS = "connect 127.0.0.1:%(adb_port)d"
56# Store the ports that vnc/adb are forwarded to, both are integers.
57ForwardedPorts = collections.namedtuple("ForwardedPorts", [constants.VNC_PORT,
58 constants.ADB_PORT])
Kevin Chengae7a49d2018-10-18 14:11:22 -070059_VNC_BIN = "ssvnc"
60_CMD_START_VNC = "%(bin)s vnc://127.0.01:%(port)d"
61_CMD_INSTALL_SSVNC = "sudo apt-get --assume-yes install ssvnc"
62_ENV_DISPLAY = "DISPLAY"
63_SSVNC_ENV_VARS = {"SSVNC_NO_ENC_WARN":"1", "SSVNC_SCALE":"auto"}
64
65_CONFIRM_CONTINUE = ("In order to display the screen to the AVD, we'll need to "
66 "install a vnc client (ssnvc). \nWould you like acloud to "
67 "install it for you? (%s) \nPress 'y' to continue or "
68 "anything else to abort it:[y] ") % _CMD_INSTALL_SSVNC
69
Fang Deng69498c32017-03-02 14:29:30 -080070
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070071class TempDir(object):
Fang Deng26e4dc12018-03-04 19:01:59 -080072 """A context manager that ceates a temporary directory.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070073
Fang Deng26e4dc12018-03-04 19:01:59 -080074 Attributes:
Sam Chiu81bdc652018-06-29 18:45:08 +080075 path: The path of the temporary directory.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070076 """
77
Fang Deng26e4dc12018-03-04 19:01:59 -080078 def __init__(self):
79 self.path = tempfile.mkdtemp()
80 os.chmod(self.path, 0o700)
81 logger.debug("Created temporary dir %s", self.path)
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070082
83 def __enter__(self):
Fang Deng26e4dc12018-03-04 19:01:59 -080084 """Enter."""
85 return self.path
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070086
Fang Deng26e4dc12018-03-04 19:01:59 -080087 def __exit__(self, exc_type, exc_value, traceback):
88 """Exit.
89
90 Args:
91 exc_type: Exception type raised within the context manager.
92 None if no execption is raised.
93 exc_value: Exception instance raised within the context manager.
94 None if no execption is raised.
95 traceback: Traceback for exeception that is raised within
96 the context manager.
97 None if no execption is raised.
98 Raises:
99 EnvironmentError or OSError when failed to delete temp directory.
100 """
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700101 try:
Fang Deng26e4dc12018-03-04 19:01:59 -0800102 if self.path:
103 shutil.rmtree(self.path)
104 logger.debug("Deleted temporary dir %s", self.path)
105 except EnvironmentError as e:
106 # Ignore error if there is no exception raised
107 # within the with-clause and the EnvironementError is
108 # about problem that directory or file does not exist.
109 if not exc_type and e.errno != errno.ENOENT:
110 raise
111 except Exception as e: # pylint: disable=W0703
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700112 if exc_type:
Fang Deng26e4dc12018-03-04 19:01:59 -0800113 logger.error(
cylan0d77ae12018-05-18 08:36:48 +0000114 "Encountered error while deleting %s: %s",
115 self.path,
116 str(e),
117 exc_info=True)
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700118 else:
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700119 raise
120
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700121
cylan0d77ae12018-05-18 08:36:48 +0000122def RetryOnException(retry_checker,
123 max_retries,
124 sleep_multiplier=0,
Fang Dengf24be082018-02-10 10:09:55 -0800125 retry_backoff_factor=1):
cylan0d77ae12018-05-18 08:36:48 +0000126 """Decorater which retries the function call if |retry_checker| returns true.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700127
cylan0d77ae12018-05-18 08:36:48 +0000128 Args:
129 retry_checker: A callback function which should take an exception instance
130 and return True if functor(*args, **kwargs) should be retried
131 when such exception is raised, and return False if it should
132 not be retried.
133 max_retries: Maximum number of retries allowed.
134 sleep_multiplier: Will sleep sleep_multiplier * attempt_count seconds if
135 retry_backoff_factor is 1. Will sleep
136 sleep_multiplier * (
137 retry_backoff_factor ** (attempt_count - 1))
138 if retry_backoff_factor != 1.
139 retry_backoff_factor: See explanation of sleep_multiplier.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700140
cylan0d77ae12018-05-18 08:36:48 +0000141 Returns:
142 The function wrapper.
143 """
144
145 def _Wrapper(func):
146 def _FunctionWrapper(*args, **kwargs):
147 return Retry(retry_checker, max_retries, func, sleep_multiplier,
148 retry_backoff_factor, *args, **kwargs)
149
150 return _FunctionWrapper
151
152 return _Wrapper
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700153
154
cylan4f73c1f2018-07-19 16:40:31 +0800155def Retry(retry_checker, max_retries, functor, sleep_multiplier,
156 retry_backoff_factor, *args, **kwargs):
cylan0d77ae12018-05-18 08:36:48 +0000157 """Conditionally retry a function.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700158
cylan0d77ae12018-05-18 08:36:48 +0000159 Args:
160 retry_checker: A callback function which should take an exception instance
161 and return True if functor(*args, **kwargs) should be retried
162 when such exception is raised, and return False if it should
163 not be retried.
164 max_retries: Maximum number of retries allowed.
165 functor: The function to call, will call functor(*args, **kwargs).
166 sleep_multiplier: Will sleep sleep_multiplier * attempt_count seconds if
167 retry_backoff_factor is 1. Will sleep
168 sleep_multiplier * (
169 retry_backoff_factor ** (attempt_count - 1))
170 if retry_backoff_factor != 1.
171 retry_backoff_factor: See explanation of sleep_multiplier.
172 *args: Arguments to pass to the functor.
173 **kwargs: Key-val based arguments to pass to the functor.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700174
cylan0d77ae12018-05-18 08:36:48 +0000175 Returns:
176 The return value of the functor.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700177
cylan0d77ae12018-05-18 08:36:48 +0000178 Raises:
179 Exception: The exception that functor(*args, **kwargs) throws.
180 """
181 attempt_count = 0
182 while attempt_count <= max_retries:
183 try:
184 attempt_count += 1
185 return_value = functor(*args, **kwargs)
186 return return_value
187 except Exception as e: # pylint: disable=W0703
188 if retry_checker(e) and attempt_count <= max_retries:
189 if retry_backoff_factor != 1:
190 sleep = sleep_multiplier * (retry_backoff_factor**
191 (attempt_count - 1))
192 else:
193 sleep = sleep_multiplier * attempt_count
Kevin Chengd25feee2018-05-24 10:15:20 -0700194 time.sleep(sleep)
cylan0d77ae12018-05-18 08:36:48 +0000195 else:
196 raise
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700197
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700198
Fang Dengf24be082018-02-10 10:09:55 -0800199def RetryExceptionType(exception_types, max_retries, functor, *args, **kwargs):
cylan0d77ae12018-05-18 08:36:48 +0000200 """Retry exception if it is one of the given types.
Fang Dengf24be082018-02-10 10:09:55 -0800201
cylan0d77ae12018-05-18 08:36:48 +0000202 Args:
203 exception_types: A tuple of exception types, e.g. (ValueError, KeyError)
204 max_retries: Max number of retries allowed.
205 functor: The function to call. Will be retried if exception is raised and
206 the exception is one of the exception_types.
207 *args: Arguments to pass to Retry function.
208 **kwargs: Key-val based arguments to pass to Retry functions.
Fang Dengf24be082018-02-10 10:09:55 -0800209
cylan0d77ae12018-05-18 08:36:48 +0000210 Returns:
211 The value returned by calling functor.
212 """
213 return Retry(lambda e: isinstance(e, exception_types), max_retries,
214 functor, *args, **kwargs)
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700215
216
217def PollAndWait(func, expected_return, timeout_exception, timeout_secs,
218 sleep_interval_secs, *args, **kwargs):
219 """Call a function until the function returns expected value or times out.
220
221 Args:
222 func: Function to call.
223 expected_return: The expected return value.
224 timeout_exception: Exception to raise when it hits timeout.
225 timeout_secs: Timeout seconds.
226 If 0 or less than zero, the function will run once and
227 we will not wait on it.
228 sleep_interval_secs: Time to sleep between two attemps.
229 *args: list of args to pass to func.
230 **kwargs: dictionary of keyword based args to pass to func.
231
232 Raises:
233 timeout_exception: if the run of function times out.
234 """
235 # TODO(fdeng): Currently this method does not kill
236 # |func|, if |func| takes longer than |timeout_secs|.
237 # We can use a more robust version from chromite.
238 start = time.time()
239 while True:
240 return_value = func(*args, **kwargs)
241 if return_value == expected_return:
242 return
243 elif time.time() - start > timeout_secs:
244 raise timeout_exception
245 else:
246 if sleep_interval_secs > 0:
247 time.sleep(sleep_interval_secs)
248
249
250def GenerateUniqueName(prefix=None, suffix=None):
Sam Chiu81bdc652018-06-29 18:45:08 +0800251 """Generate a random unique name using uuid4.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700252
253 Args:
254 prefix: String, desired prefix to prepend to the generated name.
255 suffix: String, desired suffix to append to the generated name.
256
257 Returns:
258 String, a random name.
259 """
260 name = uuid.uuid4().hex
261 if prefix:
262 name = "-".join([prefix, name])
263 if suffix:
264 name = "-".join([name, suffix])
265 return name
266
267
268def MakeTarFile(src_dict, dest):
269 """Archive files in tar.gz format to a file named as |dest|.
270
271 Args:
272 src_dict: A dictionary that maps a path to be archived
273 to the corresponding name that appears in the archive.
274 dest: String, path to output file, e.g. /tmp/myfile.tar.gz
275 """
276 logger.info("Compressing %s into %s.", src_dict.keys(), dest)
277 with tarfile.open(dest, "w:gz") as tar:
278 for src, arcname in src_dict.iteritems():
279 tar.add(src, arcname=arcname)
280
281
Fang Deng69498c32017-03-02 14:29:30 -0800282def CreateSshKeyPairIfNotExist(private_key_path, public_key_path):
283 """Create the ssh key pair if they don't exist.
284
cylan4f73c1f2018-07-19 16:40:31 +0800285 Case1. If the private key doesn't exist, we will create both the public key
286 and the private key.
287 Case2. If the private key exists but public key doesn't, we will create the
288 public key by using the private key.
289 Case3. If the public key exists but the private key doesn't, we will create
290 a new private key and overwrite the public key.
Fang Deng69498c32017-03-02 14:29:30 -0800291
292 Args:
293 private_key_path: Path to the private key file.
294 e.g. ~/.ssh/acloud_rsa
295 public_key_path: Path to the public key file.
296 e.g. ~/.ssh/acloud_rsa.pub
cylan4f73c1f2018-07-19 16:40:31 +0800297
Fang Deng69498c32017-03-02 14:29:30 -0800298 Raises:
299 error.DriverError: If failed to create the key pair.
300 """
301 public_key_path = os.path.expanduser(public_key_path)
302 private_key_path = os.path.expanduser(private_key_path)
cylan4f73c1f2018-07-19 16:40:31 +0800303 public_key_exist = os.path.exists(public_key_path)
304 private_key_exist = os.path.exists(private_key_path)
305 if public_key_exist and private_key_exist:
cylan0d77ae12018-05-18 08:36:48 +0000306 logger.debug(
cylan4f73c1f2018-07-19 16:40:31 +0800307 "The ssh private key (%s) and public key (%s) already exist,"
cylan0d77ae12018-05-18 08:36:48 +0000308 "will not automatically create the key pairs.", private_key_path,
309 public_key_path)
Fang Deng69498c32017-03-02 14:29:30 -0800310 return
cylan4f73c1f2018-07-19 16:40:31 +0800311 key_folder = os.path.dirname(private_key_path)
312 if not os.path.exists(key_folder):
313 os.makedirs(key_folder)
Fang Deng69498c32017-03-02 14:29:30 -0800314 try:
cylan4f73c1f2018-07-19 16:40:31 +0800315 if private_key_exist:
316 cmd = SSH_KEYGEN_PUB_CMD + ["-f", private_key_path]
317 with open(public_key_path, 'w') as outfile:
318 stream_content = subprocess.check_output(cmd)
319 outfile.write(
320 stream_content.rstrip('\n') + " " + getpass.getuser())
321 logger.info(
322 "The ssh public key (%s) do not exist, "
323 "automatically creating public key, calling: %s",
324 public_key_path, " ".join(cmd))
325 else:
326 cmd = SSH_KEYGEN_CMD + [
327 "-C", getpass.getuser(), "-f", private_key_path
328 ]
329 logger.info(
330 "Creating public key from private key (%s) via cmd: %s",
331 private_key_path, " ".join(cmd))
332 subprocess.check_call(cmd, stdout=sys.stderr, stderr=sys.stdout)
Fang Deng69498c32017-03-02 14:29:30 -0800333 except subprocess.CalledProcessError as e:
cylan0d77ae12018-05-18 08:36:48 +0000334 raise errors.DriverError("Failed to create ssh key pair: %s" % str(e))
Fang Deng69498c32017-03-02 14:29:30 -0800335 except OSError as e:
336 raise errors.DriverError(
cylan0d77ae12018-05-18 08:36:48 +0000337 "Failed to create ssh key pair, please make sure "
338 "'ssh-keygen' is installed: %s" % str(e))
Fang Deng69498c32017-03-02 14:29:30 -0800339
340 # By default ssh-keygen will create a public key file
341 # by append .pub to the private key file name. Rename it
342 # to what's requested by public_key_path.
343 default_pub_key_path = "%s.pub" % private_key_path
344 try:
345 if default_pub_key_path != public_key_path:
346 os.rename(default_pub_key_path, public_key_path)
347 except OSError as e:
348 raise errors.DriverError(
cylan0d77ae12018-05-18 08:36:48 +0000349 "Failed to rename %s to %s: %s" % (default_pub_key_path,
350 public_key_path, str(e)))
Fang Deng69498c32017-03-02 14:29:30 -0800351
352 logger.info("Created ssh private key (%s) and public key (%s)",
353 private_key_path, public_key_path)
354
355
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700356def VerifyRsaPubKey(rsa):
357 """Verify the format of rsa public key.
358
359 Args:
360 rsa: content of rsa public key. It should follow the format of
361 ssh-rsa AAAAB3NzaC1yc2EA.... test@test.com
362
363 Raises:
364 DriverError if the format is not correct.
365 """
366 if not rsa or not all(ord(c) < 128 for c in rsa):
367 raise errors.DriverError(
368 "rsa key is empty or contains non-ascii character: %s" % rsa)
369
370 elements = rsa.split()
371 if len(elements) != 3:
372 raise errors.DriverError("rsa key is invalid, wrong format: %s" % rsa)
373
374 key_type, data, _ = elements
375 try:
376 binary_data = base64.decodestring(data)
377 # number of bytes of int type
378 int_length = 4
379 # binary_data is like "7ssh-key..." in a binary format.
380 # The first 4 bytes should represent 7, which should be
381 # the length of the following string "ssh-key".
382 # And the next 7 bytes should be string "ssh-key".
383 # We will verify that the rsa conforms to this format.
384 # ">I" in the following line means "big-endian unsigned integer".
385 type_length = struct.unpack(">I", binary_data[:int_length])[0]
386 if binary_data[int_length:int_length + type_length] != key_type:
387 raise errors.DriverError("rsa key is invalid: %s" % rsa)
388 except (struct.error, binascii.Error) as e:
cylan0d77ae12018-05-18 08:36:48 +0000389 raise errors.DriverError(
390 "rsa key is invalid: %s, error: %s" % (rsa, str(e)))
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700391
chojoycecd004bc2018-09-13 10:39:00 +0800392def Decompress(sourcefile, dest=None):
393 """Decompress .zip or .tar.gz.
394
395 Args:
396 sourcefile: A string, a source file path to decompress.
397 dest: A string, a folder path as decompress destination.
398
399 Raises:
400 errors.UnsupportedCompressionFileType: Not supported extension.
401 """
402 logger.info("Start to decompress %s!", sourcefile)
403 dest_path = dest if dest else "."
404 if sourcefile.endswith(".tar.gz"):
405 with tarfile.open(sourcefile, "r:gz") as compressor:
406 compressor.extractall(dest_path)
407 elif sourcefile.endswith(".zip"):
408 with zipfile.ZipFile(sourcefile, 'r') as compressor:
409 compressor.extractall(dest_path)
410 else:
411 raise root_errors.UnsupportedCompressionFileType(
412 "Sorry, we could only support compression file type "
413 "for zip or tar.gz.")
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700414
Sam Chiu81bdc652018-06-29 18:45:08 +0800415# pylint: disable=old-style-class,no-init
416class TextColors:
417 """A class that defines common color ANSI code."""
418
419 HEADER = "\033[95m"
420 OKBLUE = "\033[94m"
421 OKGREEN = "\033[92m"
422 WARNING = "\033[93m"
423 FAIL = "\033[91m"
424 ENDC = "\033[0m"
425 BOLD = "\033[1m"
426 UNDERLINE = "\033[4m"
427
428
herbertxuedf01c422018-09-06 19:52:52 +0800429def PrintColorString(message, colors=TextColors.OKBLUE, **kwargs):
Sam Chiu81bdc652018-06-29 18:45:08 +0800430 """A helper function to print out colored text.
431
herbertxuedf01c422018-09-06 19:52:52 +0800432 Use print function "print(message, end="")" to show message in one line.
433 Example code:
434 DisplayMessages("Creating GCE instance...", end="")
435 # Job execute 20s
436 DisplayMessages("Done! (20s)")
437 Display:
438 Creating GCE instance...
439 # After job finished, messages update as following:
440 Creating GCE instance...Done! (20s)
441
Sam Chiu81bdc652018-06-29 18:45:08 +0800442 Args:
443 message: String, the message text.
444 colors: String, color code.
herbertxuedf01c422018-09-06 19:52:52 +0800445 **kwargs: dictionary of keyword based args to pass to func.
Sam Chiu81bdc652018-06-29 18:45:08 +0800446 """
herbertxuedf01c422018-09-06 19:52:52 +0800447 print(colors + message + TextColors.ENDC, **kwargs)
448 sys.stdout.flush()
Sam Chiu81bdc652018-06-29 18:45:08 +0800449
450
451def InteractWithQuestion(question, colors=TextColors.WARNING):
452 """A helper function to define the common way to run interactive cmd.
453
454 Args:
455 question: String, the question to ask user.
456 colors: String, color code.
457
458 Returns:
459 String, input from user.
460 """
461 return str(raw_input(colors + question + TextColors.ENDC).strip())
462
herbertxuedf01c422018-09-06 19:52:52 +0800463
herbertxue34776bb2018-07-03 21:57:48 +0800464def GetUserAnswerYes(question):
465 """Ask user about acloud setup question.
466
467 Args:
468 question: String, ask question for user.
469 Ex: "Are you sure to change bucket name:[y/n]"
470
471 Returns:
472 Boolean, True if answer is "Yes", False otherwise.
473 """
474 answer = InteractWithQuestion(question)
475 return answer.lower() in constants.USER_ANSWER_YES
476
Sam Chiu81bdc652018-06-29 18:45:08 +0800477
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700478class BatchHttpRequestExecutor(object):
479 """A helper class that executes requests in batch with retry.
480
481 This executor executes http requests in a batch and retry
482 those that have failed. It iteratively updates the dictionary
483 self._final_results with latest results, which can be retrieved
484 via GetResults.
485 """
486
487 def __init__(self,
488 execute_once_functor,
489 requests,
490 retry_http_codes=None,
491 max_retry=None,
492 sleep=None,
493 backoff_factor=None,
494 other_retriable_errors=None):
495 """Initializes the executor.
496
497 Args:
498 execute_once_functor: A function that execute requests in batch once.
499 It should return a dictionary like
500 {request_id: (response, exception)}
501 requests: A dictionary where key is request id picked by caller,
502 and value is a apiclient.http.HttpRequest.
503 retry_http_codes: A list of http codes to retry.
Fang Dengf24be082018-02-10 10:09:55 -0800504 max_retry: See utils.Retry.
505 sleep: See utils.Retry.
506 backoff_factor: See utils.Retry.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700507 other_retriable_errors: A tuple of error types that should be retried
508 other than errors.HttpError.
509 """
510 self._execute_once_functor = execute_once_functor
511 self._requests = requests
512 # A dictionary that maps request id to pending request.
513 self._pending_requests = {}
514 # A dictionary that maps request id to a tuple (response, exception).
515 self._final_results = {}
516 self._retry_http_codes = retry_http_codes
517 self._max_retry = max_retry
518 self._sleep = sleep
519 self._backoff_factor = backoff_factor
520 self._other_retriable_errors = other_retriable_errors
521
522 def _ShoudRetry(self, exception):
Sam Chiu81bdc652018-06-29 18:45:08 +0800523 """Check if an exception is retriable.
524
525 Args:
526 exception: An exception instance.
527 """
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700528 if isinstance(exception, self._other_retriable_errors):
529 return True
530
cylan0d77ae12018-05-18 08:36:48 +0000531 if (isinstance(exception, errors.HttpError)
532 and exception.code in self._retry_http_codes):
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700533 return True
534 return False
535
536 def _ExecuteOnce(self):
537 """Executes pending requests and update it with failed, retriable ones.
538
539 Raises:
540 HasRetriableRequestsError: if some requests fail and are retriable.
541 """
542 results = self._execute_once_functor(self._pending_requests)
543 # Update final_results with latest results.
544 self._final_results.update(results)
545 # Clear pending_requests
546 self._pending_requests.clear()
547 for request_id, result in results.iteritems():
548 exception = result[1]
549 if exception is not None and self._ShoudRetry(exception):
550 # If this is a retriable exception, put it in pending_requests
551 self._pending_requests[request_id] = self._requests[request_id]
552 if self._pending_requests:
553 # If there is still retriable requests pending, raise an error
Fang Dengf24be082018-02-10 10:09:55 -0800554 # so that Retry will retry this function with pending_requests.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700555 raise errors.HasRetriableRequestsError(
cylan0d77ae12018-05-18 08:36:48 +0000556 "Retriable errors: %s" %
557 [str(results[rid][1]) for rid in self._pending_requests])
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700558
559 def Execute(self):
560 """Executes the requests and retry if necessary.
561
562 Will populate self._final_results.
563 """
cylan0d77ae12018-05-18 08:36:48 +0000564
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700565 def _ShouldRetryHandler(exc):
566 """Check if |exc| is a retriable exception.
567
568 Args:
569 exc: An exception.
570
571 Returns:
572 True if exception is of type HasRetriableRequestsError; False otherwise.
573 """
574 should_retry = isinstance(exc, errors.HasRetriableRequestsError)
575 if should_retry:
576 logger.info("Will retry failed requests.", exc_info=True)
577 logger.info("%s", exc)
578 return should_retry
579
580 try:
581 self._pending_requests = self._requests.copy()
Fang Dengf24be082018-02-10 10:09:55 -0800582 Retry(
cylan0d77ae12018-05-18 08:36:48 +0000583 _ShouldRetryHandler,
584 max_retries=self._max_retry,
Fang Dengf24be082018-02-10 10:09:55 -0800585 functor=self._ExecuteOnce,
586 sleep_multiplier=self._sleep,
587 retry_backoff_factor=self._backoff_factor)
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700588 except errors.HasRetriableRequestsError:
589 logger.debug("Some requests did not succeed after retry.")
590
591 def GetResults(self):
592 """Returns final results.
593
594 Returns:
595 results, a dictionary in the following format
596 {request_id: (response, exception)}
597 request_ids are those from requests; response
598 is the http response for the request or None on error;
599 exception is an instance of DriverError or None if no error.
600 """
601 return self._final_results
cylan31fc5332018-09-17 22:12:08 +0800602
603
604class TimeExecute(object):
605 """Count the function execute time."""
606
607 def __init__(self, function_description=None, print_before_call=True, print_status=True):
608 """Initializes the class.
609
610 Args:
611 function_description: String that describes function (e.g."Creating
612 Instance...")
613 print_before_call: Boolean, print the function description before
614 calling the function, default True.
615 print_status: Boolean, print the status of the function after the
616 function has completed, default True ("OK" or "Fail").
617 """
618 self._function_description = function_description
619 self._print_before_call = print_before_call
620 self._print_status = print_status
621
622 def __call__(self, func):
623 def DecoratorFunction(*args, **kargs):
624 """Decorator function.
625
626 Args:
627 *args: Arguments to pass to the functor.
628 **kwargs: Key-val based arguments to pass to the functor.
629
630 Raises:
631 Exception: The exception that functor(*args, **kwargs) throws.
632 """
633 timestart = time.time()
634 if self._print_before_call:
635 PrintColorString("%s ..."% self._function_description, end="")
636 try:
637 result = func(*args, **kargs)
638 if not self._print_before_call:
639 PrintColorString("%s (%ds)" % (self._function_description,
640 time.time()-timestart),
641 TextColors.OKGREEN)
642 if self._print_status:
643 PrintColorString("OK! (%ds)" % (time.time()-timestart),
644 TextColors.OKGREEN)
645 return result
646 except:
647 if self._print_status:
648 PrintColorString("Fail! (%ds)" % (time.time()-timestart),
649 TextColors.FAIL)
650 raise
651 return DecoratorFunction
cylan66713722018-10-06 01:38:26 +0800652
653
654def PickFreePort():
655 """Helper to pick a free port.
656
657 Returns:
658 Integer, a free port number.
659 """
660 tcp_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
661 tcp_socket.bind(("", 0))
662 port = tcp_socket.getsockname()[1]
663 tcp_socket.close()
664 return port
665
666
667def _ExecuteCommand(cmd, args):
668 """Execute command.
669
670 Args:
671 cmd: Strings of execute binary name.
672 args: List of args to pass in with cmd.
673
674 Raises:
675 errors.NoExecuteBin: Can't find the execute bin file.
676 """
677 bin_path = find_executable(cmd)
678 if not bin_path:
679 raise root_errors.NoExecuteCmd("unable to locate %s" % cmd)
680 command = [bin_path] + args
681 logger.debug("Running '%s'", ' '.join(command))
682 with open(os.devnull, "w") as dev_null:
683 subprocess.check_call(command, stderr=dev_null, stdout=dev_null)
684
685
686# pylint: disable=too-many-locals
687def AutoConnect(ip_addr, rsa_key_file, target_vnc_port, target_adb_port, ssh_user):
688 """Autoconnect to an AVD instance.
689
690 Args:
691 ip_addr: String, use to build the adb & vnc tunnel between local
692 and remote instance.
693 rsa_key_file: String, Private key file path to use when creating
694 the ssh tunnels.
695 target_vnc_port: Integer of target vnc port number.
696 target_adb_port: Integer of target adb port number.
697 ssh_user: String of user login into the instance.
698
699 Returns:
700 NamedTuple of (vnc_port, adb_port) SSHTUNNEL of the connect, both are
701 integers.
702 """
703 local_free_vnc_port = PickFreePort()
704 local_free_adb_port = PickFreePort()
705 try:
706 ssh_tunnel_args = _SSH_TUNNEL_ARGS % {
707 "rsa_key_file": rsa_key_file,
708 "vnc_port": local_free_vnc_port,
709 "adb_port": local_free_adb_port,
710 "target_vnc_port": target_vnc_port,
711 "target_adb_port": target_adb_port,
712 "ssh_user": ssh_user,
713 "ip_addr": ip_addr}
Kevin Cheng835a4152018-10-11 10:46:57 -0700714 _ExecuteCommand(constants.SSH_BIN, ssh_tunnel_args.split())
cylan66713722018-10-06 01:38:26 +0800715 except subprocess.CalledProcessError:
716 PrintColorString("Failed to create ssh tunnels, retry with '#acloud "
717 "reconnect'.", TextColors.FAIL)
718 try:
719 adb_connect_args = _ADB_CONNECT_ARGS % {"adb_port": local_free_adb_port}
Kevin Cheng835a4152018-10-11 10:46:57 -0700720 _ExecuteCommand(constants.ADB_BIN, adb_connect_args.split())
cylan66713722018-10-06 01:38:26 +0800721 except subprocess.CalledProcessError:
722 PrintColorString("Failed to adb connect, retry with "
723 "'#acloud reconnect'", TextColors.FAIL)
724
725 return ForwardedPorts(vnc_port=local_free_vnc_port,
726 adb_port=local_free_adb_port)
Kevin Chengeb85e862018-10-09 15:35:13 -0700727
728
729def GetAnswerFromList(answer_list, enable_choose_all=False):
730 """Get answer from a list.
731
732 Args:
733 answer_list: list of the answers to choose from.
734
735 Return:
736 List holding the answer(s).
737 """
738 print("[0] to exit.")
739 start_index = 1
740 if enable_choose_all:
741 start_index = 2
742 print("[1] for all.")
743 for num, item in enumerate(answer_list, start_index):
744 print("[%d] %s" % (num, item))
745
746 choice = -1
747 max_choice = len(answer_list) + 1
748 while True:
749 try:
750 choice = raw_input("Enter your choice[0-%d]: " % max_choice)
751 choice = int(choice)
752 except ValueError:
753 print("'%s' is not a valid integer.", choice)
754 continue
755 # Filter out choices
756 if choice == 0:
757 print("Exiting acloud.")
758 sys.exit()
759 if enable_choose_all and choice == 1:
760 return answer_list
761 if choice < 0 or choice > max_choice:
762 print("please choose between 0 and %d" % max_choice)
763 else:
764 return [answer_list[choice-start_index]]
Kevin Chengae7a49d2018-10-18 14:11:22 -0700765
766
767def LaunchVncClient(port=constants.DEFAULT_VNC_PORT):
768 """Launch ssvnc.
769
770 Args:
771 port: Integer, port number.
772 """
773 try:
774 os.environ[_ENV_DISPLAY]
775 except KeyError:
776 PrintColorString("Remote terminal can't support VNC. "
777 "Skipping VNC startup.", TextColors.FAIL)
778 return
779
780 if not find_executable(_VNC_BIN):
781 if GetUserAnswerYes(_CONFIRM_CONTINUE):
782 try:
783 PrintColorString("Installing ssvnc vnc client... ", end="")
784 sys.stdout.flush()
785 subprocess.check_output(_CMD_INSTALL_SSVNC, shell=True)
786 PrintColorString("Done", TextColors.OKGREEN)
787 except subprocess.CalledProcessError as cpe:
788 PrintColorString("Failed to install ssvnc: %s" %
789 cpe.output, TextColors.FAIL)
790 return
791 else:
792 return
793 ssvnc_env = os.environ.copy()
794 ssvnc_env.update(_SSVNC_ENV_VARS)
795 ssvnc_args = _CMD_START_VNC % {"bin":find_executable(_VNC_BIN),
796 "port":port}
797 subprocess.Popen(ssvnc_args.split(), env=ssvnc_env)
798
799
800def PrintDeviceSummary(report):
801 """Display summary of devices created.
802
803 -Display created device details from the report instance.
804 report example:
805 'data': [{'devices':[{'instance_name': 'ins-f6a397-none-53363',
806 'ip': u'35.234.10.162'}]}]
807 -Display error message from report.error.
808
809 Args:
810 report: A Report instance.
811 """
812 PrintColorString("\n")
813 PrintColorString("Device(s) created:")
814 for device in report.data.get("devices", []):
815 adb_serial = "(None)"
816 adb_port = device.get("adb_port")
817 if adb_port:
818 adb_serial = constants.LOCALHOST_ADB_SERIAL % adb_port
819 instance_name = device.get("instance_name")
820 instance_ip = device.get("ip")
821 instance_details = "" if not instance_name else "(%s[%s])" % (
822 instance_name, instance_ip)
823 PrintColorString(" - device serial: %s %s" % (adb_serial,
824 instance_details))
825
826 # TODO(b/117245508): Help user to delete instance if it got created.
827 if report.errors:
828 error_msg = "\n".join(report.errors)
829 PrintColorString("Fail in:\n%s\n" % error_msg, TextColors.FAIL)