blob: 8f39455c532f2b8fea2934f16e9b64e95d662a2f [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
herbertxue07293a32018-11-05 20:40:11 +080026import grp
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070027import logging
28import os
29import shutil
30import struct
cylan66713722018-10-06 01:38:26 +080031import socket
Fang Deng69498c32017-03-02 14:29:30 -080032import subprocess
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070033import sys
34import tarfile
35import tempfile
36import time
37import 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"
herbertxue07293a32018-11-05 20:40:11 +080061_CMD_PGREP = "pgrep"
62_CMD_SG = "sg "
Kevin Chengae7a49d2018-10-18 14:11:22 -070063_CMD_START_VNC = "%(bin)s vnc://127.0.01:%(port)d"
64_CMD_INSTALL_SSVNC = "sudo apt-get --assume-yes install ssvnc"
65_ENV_DISPLAY = "DISPLAY"
Sam Chiu7a477f52018-10-22 11:20:36 +080066_SSVNC_ENV_VARS = {"SSVNC_NO_ENC_WARN": "1", "SSVNC_SCALE": "auto"}
Kevin Cheng53aa5a52018-12-03 01:33:55 -080067_DEFAULT_DISPLAY_SCALE = 1.0
Kevin Chengae7a49d2018-10-18 14:11:22 -070068
69_CONFIRM_CONTINUE = ("In order to display the screen to the AVD, we'll need to "
70 "install a vnc client (ssnvc). \nWould you like acloud to "
71 "install it for you? (%s) \nPress 'y' to continue or "
72 "anything else to abort it:[y] ") % _CMD_INSTALL_SSVNC
73
Fang Deng69498c32017-03-02 14:29:30 -080074
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070075class TempDir(object):
Fang Deng26e4dc12018-03-04 19:01:59 -080076 """A context manager that ceates a temporary directory.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070077
Fang Deng26e4dc12018-03-04 19:01:59 -080078 Attributes:
Sam Chiu81bdc652018-06-29 18:45:08 +080079 path: The path of the temporary directory.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070080 """
81
Fang Deng26e4dc12018-03-04 19:01:59 -080082 def __init__(self):
83 self.path = tempfile.mkdtemp()
84 os.chmod(self.path, 0o700)
85 logger.debug("Created temporary dir %s", self.path)
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070086
87 def __enter__(self):
Fang Deng26e4dc12018-03-04 19:01:59 -080088 """Enter."""
89 return self.path
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070090
Fang Deng26e4dc12018-03-04 19:01:59 -080091 def __exit__(self, exc_type, exc_value, traceback):
92 """Exit.
93
94 Args:
95 exc_type: Exception type raised within the context manager.
96 None if no execption is raised.
97 exc_value: Exception instance raised within the context manager.
98 None if no execption is raised.
99 traceback: Traceback for exeception that is raised within
100 the context manager.
101 None if no execption is raised.
102 Raises:
103 EnvironmentError or OSError when failed to delete temp directory.
104 """
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700105 try:
Fang Deng26e4dc12018-03-04 19:01:59 -0800106 if self.path:
107 shutil.rmtree(self.path)
108 logger.debug("Deleted temporary dir %s", self.path)
109 except EnvironmentError as e:
110 # Ignore error if there is no exception raised
111 # within the with-clause and the EnvironementError is
112 # about problem that directory or file does not exist.
113 if not exc_type and e.errno != errno.ENOENT:
114 raise
115 except Exception as e: # pylint: disable=W0703
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700116 if exc_type:
Fang Deng26e4dc12018-03-04 19:01:59 -0800117 logger.error(
cylan0d77ae12018-05-18 08:36:48 +0000118 "Encountered error while deleting %s: %s",
119 self.path,
120 str(e),
121 exc_info=True)
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700122 else:
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700123 raise
124
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700125
cylan0d77ae12018-05-18 08:36:48 +0000126def RetryOnException(retry_checker,
127 max_retries,
128 sleep_multiplier=0,
Fang Dengf24be082018-02-10 10:09:55 -0800129 retry_backoff_factor=1):
cylan0d77ae12018-05-18 08:36:48 +0000130 """Decorater which retries the function call if |retry_checker| returns true.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700131
cylan0d77ae12018-05-18 08:36:48 +0000132 Args:
133 retry_checker: A callback function which should take an exception instance
134 and return True if functor(*args, **kwargs) should be retried
135 when such exception is raised, and return False if it should
136 not be retried.
137 max_retries: Maximum number of retries allowed.
138 sleep_multiplier: Will sleep sleep_multiplier * attempt_count seconds if
139 retry_backoff_factor is 1. Will sleep
140 sleep_multiplier * (
141 retry_backoff_factor ** (attempt_count - 1))
142 if retry_backoff_factor != 1.
143 retry_backoff_factor: See explanation of sleep_multiplier.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700144
cylan0d77ae12018-05-18 08:36:48 +0000145 Returns:
146 The function wrapper.
147 """
148
149 def _Wrapper(func):
150 def _FunctionWrapper(*args, **kwargs):
151 return Retry(retry_checker, max_retries, func, sleep_multiplier,
152 retry_backoff_factor, *args, **kwargs)
153
154 return _FunctionWrapper
155
156 return _Wrapper
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700157
158
cylan4f73c1f2018-07-19 16:40:31 +0800159def Retry(retry_checker, max_retries, functor, sleep_multiplier,
160 retry_backoff_factor, *args, **kwargs):
cylan0d77ae12018-05-18 08:36:48 +0000161 """Conditionally retry a function.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700162
cylan0d77ae12018-05-18 08:36:48 +0000163 Args:
164 retry_checker: A callback function which should take an exception instance
165 and return True if functor(*args, **kwargs) should be retried
166 when such exception is raised, and return False if it should
167 not be retried.
168 max_retries: Maximum number of retries allowed.
169 functor: The function to call, will call functor(*args, **kwargs).
170 sleep_multiplier: Will sleep sleep_multiplier * attempt_count seconds if
171 retry_backoff_factor is 1. Will sleep
172 sleep_multiplier * (
173 retry_backoff_factor ** (attempt_count - 1))
174 if retry_backoff_factor != 1.
175 retry_backoff_factor: See explanation of sleep_multiplier.
176 *args: Arguments to pass to the functor.
177 **kwargs: Key-val based arguments to pass to the functor.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700178
cylan0d77ae12018-05-18 08:36:48 +0000179 Returns:
180 The return value of the functor.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700181
cylan0d77ae12018-05-18 08:36:48 +0000182 Raises:
183 Exception: The exception that functor(*args, **kwargs) throws.
184 """
185 attempt_count = 0
186 while attempt_count <= max_retries:
187 try:
188 attempt_count += 1
189 return_value = functor(*args, **kwargs)
190 return return_value
191 except Exception as e: # pylint: disable=W0703
192 if retry_checker(e) and attempt_count <= max_retries:
193 if retry_backoff_factor != 1:
194 sleep = sleep_multiplier * (retry_backoff_factor**
195 (attempt_count - 1))
196 else:
197 sleep = sleep_multiplier * attempt_count
Kevin Chengd25feee2018-05-24 10:15:20 -0700198 time.sleep(sleep)
cylan0d77ae12018-05-18 08:36:48 +0000199 else:
200 raise
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700201
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700202
Fang Dengf24be082018-02-10 10:09:55 -0800203def RetryExceptionType(exception_types, max_retries, functor, *args, **kwargs):
cylan0d77ae12018-05-18 08:36:48 +0000204 """Retry exception if it is one of the given types.
Fang Dengf24be082018-02-10 10:09:55 -0800205
cylan0d77ae12018-05-18 08:36:48 +0000206 Args:
207 exception_types: A tuple of exception types, e.g. (ValueError, KeyError)
208 max_retries: Max number of retries allowed.
209 functor: The function to call. Will be retried if exception is raised and
210 the exception is one of the exception_types.
211 *args: Arguments to pass to Retry function.
212 **kwargs: Key-val based arguments to pass to Retry functions.
Fang Dengf24be082018-02-10 10:09:55 -0800213
cylan0d77ae12018-05-18 08:36:48 +0000214 Returns:
215 The value returned by calling functor.
216 """
217 return Retry(lambda e: isinstance(e, exception_types), max_retries,
218 functor, *args, **kwargs)
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700219
220
221def PollAndWait(func, expected_return, timeout_exception, timeout_secs,
222 sleep_interval_secs, *args, **kwargs):
223 """Call a function until the function returns expected value or times out.
224
225 Args:
226 func: Function to call.
227 expected_return: The expected return value.
228 timeout_exception: Exception to raise when it hits timeout.
229 timeout_secs: Timeout seconds.
230 If 0 or less than zero, the function will run once and
231 we will not wait on it.
232 sleep_interval_secs: Time to sleep between two attemps.
233 *args: list of args to pass to func.
234 **kwargs: dictionary of keyword based args to pass to func.
235
236 Raises:
237 timeout_exception: if the run of function times out.
238 """
239 # TODO(fdeng): Currently this method does not kill
240 # |func|, if |func| takes longer than |timeout_secs|.
241 # We can use a more robust version from chromite.
242 start = time.time()
243 while True:
244 return_value = func(*args, **kwargs)
245 if return_value == expected_return:
246 return
247 elif time.time() - start > timeout_secs:
248 raise timeout_exception
249 else:
250 if sleep_interval_secs > 0:
251 time.sleep(sleep_interval_secs)
252
253
254def GenerateUniqueName(prefix=None, suffix=None):
Sam Chiu81bdc652018-06-29 18:45:08 +0800255 """Generate a random unique name using uuid4.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700256
257 Args:
258 prefix: String, desired prefix to prepend to the generated name.
259 suffix: String, desired suffix to append to the generated name.
260
261 Returns:
262 String, a random name.
263 """
264 name = uuid.uuid4().hex
265 if prefix:
266 name = "-".join([prefix, name])
267 if suffix:
268 name = "-".join([name, suffix])
269 return name
270
271
272def MakeTarFile(src_dict, dest):
273 """Archive files in tar.gz format to a file named as |dest|.
274
275 Args:
276 src_dict: A dictionary that maps a path to be archived
277 to the corresponding name that appears in the archive.
278 dest: String, path to output file, e.g. /tmp/myfile.tar.gz
279 """
280 logger.info("Compressing %s into %s.", src_dict.keys(), dest)
281 with tarfile.open(dest, "w:gz") as tar:
282 for src, arcname in src_dict.iteritems():
283 tar.add(src, arcname=arcname)
284
285
Fang Deng69498c32017-03-02 14:29:30 -0800286def CreateSshKeyPairIfNotExist(private_key_path, public_key_path):
287 """Create the ssh key pair if they don't exist.
288
cylan4f73c1f2018-07-19 16:40:31 +0800289 Case1. If the private key doesn't exist, we will create both the public key
290 and the private key.
291 Case2. If the private key exists but public key doesn't, we will create the
292 public key by using the private key.
293 Case3. If the public key exists but the private key doesn't, we will create
294 a new private key and overwrite the public key.
Fang Deng69498c32017-03-02 14:29:30 -0800295
296 Args:
297 private_key_path: Path to the private key file.
298 e.g. ~/.ssh/acloud_rsa
299 public_key_path: Path to the public key file.
300 e.g. ~/.ssh/acloud_rsa.pub
cylan4f73c1f2018-07-19 16:40:31 +0800301
Fang Deng69498c32017-03-02 14:29:30 -0800302 Raises:
303 error.DriverError: If failed to create the key pair.
304 """
305 public_key_path = os.path.expanduser(public_key_path)
306 private_key_path = os.path.expanduser(private_key_path)
cylan4f73c1f2018-07-19 16:40:31 +0800307 public_key_exist = os.path.exists(public_key_path)
308 private_key_exist = os.path.exists(private_key_path)
309 if public_key_exist and private_key_exist:
cylan0d77ae12018-05-18 08:36:48 +0000310 logger.debug(
cylan4f73c1f2018-07-19 16:40:31 +0800311 "The ssh private key (%s) and public key (%s) already exist,"
cylan0d77ae12018-05-18 08:36:48 +0000312 "will not automatically create the key pairs.", private_key_path,
313 public_key_path)
Fang Deng69498c32017-03-02 14:29:30 -0800314 return
cylan4f73c1f2018-07-19 16:40:31 +0800315 key_folder = os.path.dirname(private_key_path)
316 if not os.path.exists(key_folder):
317 os.makedirs(key_folder)
Fang Deng69498c32017-03-02 14:29:30 -0800318 try:
cylan4f73c1f2018-07-19 16:40:31 +0800319 if private_key_exist:
320 cmd = SSH_KEYGEN_PUB_CMD + ["-f", private_key_path]
321 with open(public_key_path, 'w') as outfile:
322 stream_content = subprocess.check_output(cmd)
323 outfile.write(
324 stream_content.rstrip('\n') + " " + getpass.getuser())
325 logger.info(
326 "The ssh public key (%s) do not exist, "
327 "automatically creating public key, calling: %s",
328 public_key_path, " ".join(cmd))
329 else:
330 cmd = SSH_KEYGEN_CMD + [
331 "-C", getpass.getuser(), "-f", private_key_path
332 ]
333 logger.info(
334 "Creating public key from private key (%s) via cmd: %s",
335 private_key_path, " ".join(cmd))
336 subprocess.check_call(cmd, stdout=sys.stderr, stderr=sys.stdout)
Fang Deng69498c32017-03-02 14:29:30 -0800337 except subprocess.CalledProcessError as e:
cylan0d77ae12018-05-18 08:36:48 +0000338 raise errors.DriverError("Failed to create ssh key pair: %s" % str(e))
Fang Deng69498c32017-03-02 14:29:30 -0800339 except OSError as e:
340 raise errors.DriverError(
cylan0d77ae12018-05-18 08:36:48 +0000341 "Failed to create ssh key pair, please make sure "
342 "'ssh-keygen' is installed: %s" % str(e))
Fang Deng69498c32017-03-02 14:29:30 -0800343
344 # By default ssh-keygen will create a public key file
345 # by append .pub to the private key file name. Rename it
346 # to what's requested by public_key_path.
347 default_pub_key_path = "%s.pub" % private_key_path
348 try:
349 if default_pub_key_path != public_key_path:
350 os.rename(default_pub_key_path, public_key_path)
351 except OSError as e:
352 raise errors.DriverError(
cylan0d77ae12018-05-18 08:36:48 +0000353 "Failed to rename %s to %s: %s" % (default_pub_key_path,
354 public_key_path, str(e)))
Fang Deng69498c32017-03-02 14:29:30 -0800355
356 logger.info("Created ssh private key (%s) and public key (%s)",
357 private_key_path, public_key_path)
358
359
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700360def VerifyRsaPubKey(rsa):
361 """Verify the format of rsa public key.
362
363 Args:
364 rsa: content of rsa public key. It should follow the format of
365 ssh-rsa AAAAB3NzaC1yc2EA.... test@test.com
366
367 Raises:
368 DriverError if the format is not correct.
369 """
370 if not rsa or not all(ord(c) < 128 for c in rsa):
371 raise errors.DriverError(
372 "rsa key is empty or contains non-ascii character: %s" % rsa)
373
374 elements = rsa.split()
375 if len(elements) != 3:
376 raise errors.DriverError("rsa key is invalid, wrong format: %s" % rsa)
377
378 key_type, data, _ = elements
379 try:
380 binary_data = base64.decodestring(data)
381 # number of bytes of int type
382 int_length = 4
383 # binary_data is like "7ssh-key..." in a binary format.
384 # The first 4 bytes should represent 7, which should be
385 # the length of the following string "ssh-key".
386 # And the next 7 bytes should be string "ssh-key".
387 # We will verify that the rsa conforms to this format.
388 # ">I" in the following line means "big-endian unsigned integer".
389 type_length = struct.unpack(">I", binary_data[:int_length])[0]
390 if binary_data[int_length:int_length + type_length] != key_type:
391 raise errors.DriverError("rsa key is invalid: %s" % rsa)
392 except (struct.error, binascii.Error) as e:
cylan0d77ae12018-05-18 08:36:48 +0000393 raise errors.DriverError(
394 "rsa key is invalid: %s, error: %s" % (rsa, str(e)))
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700395
chojoycecd004bc2018-09-13 10:39:00 +0800396def Decompress(sourcefile, dest=None):
397 """Decompress .zip or .tar.gz.
398
399 Args:
400 sourcefile: A string, a source file path to decompress.
401 dest: A string, a folder path as decompress destination.
402
403 Raises:
404 errors.UnsupportedCompressionFileType: Not supported extension.
405 """
406 logger.info("Start to decompress %s!", sourcefile)
407 dest_path = dest if dest else "."
408 if sourcefile.endswith(".tar.gz"):
409 with tarfile.open(sourcefile, "r:gz") as compressor:
410 compressor.extractall(dest_path)
411 elif sourcefile.endswith(".zip"):
412 with zipfile.ZipFile(sourcefile, 'r') as compressor:
413 compressor.extractall(dest_path)
414 else:
415 raise root_errors.UnsupportedCompressionFileType(
416 "Sorry, we could only support compression file type "
417 "for zip or tar.gz.")
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700418
Sam Chiu81bdc652018-06-29 18:45:08 +0800419# pylint: disable=old-style-class,no-init
420class TextColors:
421 """A class that defines common color ANSI code."""
422
423 HEADER = "\033[95m"
424 OKBLUE = "\033[94m"
425 OKGREEN = "\033[92m"
426 WARNING = "\033[93m"
427 FAIL = "\033[91m"
428 ENDC = "\033[0m"
429 BOLD = "\033[1m"
430 UNDERLINE = "\033[4m"
431
432
herbertxuedf01c422018-09-06 19:52:52 +0800433def PrintColorString(message, colors=TextColors.OKBLUE, **kwargs):
Sam Chiu81bdc652018-06-29 18:45:08 +0800434 """A helper function to print out colored text.
435
herbertxuedf01c422018-09-06 19:52:52 +0800436 Use print function "print(message, end="")" to show message in one line.
437 Example code:
438 DisplayMessages("Creating GCE instance...", end="")
439 # Job execute 20s
440 DisplayMessages("Done! (20s)")
441 Display:
442 Creating GCE instance...
443 # After job finished, messages update as following:
444 Creating GCE instance...Done! (20s)
445
Sam Chiu81bdc652018-06-29 18:45:08 +0800446 Args:
447 message: String, the message text.
448 colors: String, color code.
herbertxuedf01c422018-09-06 19:52:52 +0800449 **kwargs: dictionary of keyword based args to pass to func.
Sam Chiu81bdc652018-06-29 18:45:08 +0800450 """
herbertxuedf01c422018-09-06 19:52:52 +0800451 print(colors + message + TextColors.ENDC, **kwargs)
452 sys.stdout.flush()
Sam Chiu81bdc652018-06-29 18:45:08 +0800453
454
455def InteractWithQuestion(question, colors=TextColors.WARNING):
456 """A helper function to define the common way to run interactive cmd.
457
458 Args:
459 question: String, the question to ask user.
460 colors: String, color code.
461
462 Returns:
463 String, input from user.
464 """
465 return str(raw_input(colors + question + TextColors.ENDC).strip())
466
herbertxuedf01c422018-09-06 19:52:52 +0800467
herbertxue34776bb2018-07-03 21:57:48 +0800468def GetUserAnswerYes(question):
469 """Ask user about acloud setup question.
470
471 Args:
472 question: String, ask question for user.
473 Ex: "Are you sure to change bucket name:[y/n]"
474
475 Returns:
476 Boolean, True if answer is "Yes", False otherwise.
477 """
478 answer = InteractWithQuestion(question)
479 return answer.lower() in constants.USER_ANSWER_YES
480
Sam Chiu81bdc652018-06-29 18:45:08 +0800481
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700482class BatchHttpRequestExecutor(object):
483 """A helper class that executes requests in batch with retry.
484
485 This executor executes http requests in a batch and retry
486 those that have failed. It iteratively updates the dictionary
487 self._final_results with latest results, which can be retrieved
488 via GetResults.
489 """
490
491 def __init__(self,
492 execute_once_functor,
493 requests,
494 retry_http_codes=None,
495 max_retry=None,
496 sleep=None,
497 backoff_factor=None,
498 other_retriable_errors=None):
499 """Initializes the executor.
500
501 Args:
502 execute_once_functor: A function that execute requests in batch once.
503 It should return a dictionary like
504 {request_id: (response, exception)}
505 requests: A dictionary where key is request id picked by caller,
506 and value is a apiclient.http.HttpRequest.
507 retry_http_codes: A list of http codes to retry.
Fang Dengf24be082018-02-10 10:09:55 -0800508 max_retry: See utils.Retry.
509 sleep: See utils.Retry.
510 backoff_factor: See utils.Retry.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700511 other_retriable_errors: A tuple of error types that should be retried
512 other than errors.HttpError.
513 """
514 self._execute_once_functor = execute_once_functor
515 self._requests = requests
516 # A dictionary that maps request id to pending request.
517 self._pending_requests = {}
518 # A dictionary that maps request id to a tuple (response, exception).
519 self._final_results = {}
520 self._retry_http_codes = retry_http_codes
521 self._max_retry = max_retry
522 self._sleep = sleep
523 self._backoff_factor = backoff_factor
524 self._other_retriable_errors = other_retriable_errors
525
526 def _ShoudRetry(self, exception):
Sam Chiu81bdc652018-06-29 18:45:08 +0800527 """Check if an exception is retriable.
528
529 Args:
530 exception: An exception instance.
531 """
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700532 if isinstance(exception, self._other_retriable_errors):
533 return True
534
cylan0d77ae12018-05-18 08:36:48 +0000535 if (isinstance(exception, errors.HttpError)
536 and exception.code in self._retry_http_codes):
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700537 return True
538 return False
539
540 def _ExecuteOnce(self):
541 """Executes pending requests and update it with failed, retriable ones.
542
543 Raises:
544 HasRetriableRequestsError: if some requests fail and are retriable.
545 """
546 results = self._execute_once_functor(self._pending_requests)
547 # Update final_results with latest results.
548 self._final_results.update(results)
549 # Clear pending_requests
550 self._pending_requests.clear()
551 for request_id, result in results.iteritems():
552 exception = result[1]
553 if exception is not None and self._ShoudRetry(exception):
554 # If this is a retriable exception, put it in pending_requests
555 self._pending_requests[request_id] = self._requests[request_id]
556 if self._pending_requests:
557 # If there is still retriable requests pending, raise an error
Fang Dengf24be082018-02-10 10:09:55 -0800558 # so that Retry will retry this function with pending_requests.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700559 raise errors.HasRetriableRequestsError(
cylan0d77ae12018-05-18 08:36:48 +0000560 "Retriable errors: %s" %
561 [str(results[rid][1]) for rid in self._pending_requests])
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700562
563 def Execute(self):
564 """Executes the requests and retry if necessary.
565
566 Will populate self._final_results.
567 """
cylan0d77ae12018-05-18 08:36:48 +0000568
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700569 def _ShouldRetryHandler(exc):
570 """Check if |exc| is a retriable exception.
571
572 Args:
573 exc: An exception.
574
575 Returns:
576 True if exception is of type HasRetriableRequestsError; False otherwise.
577 """
578 should_retry = isinstance(exc, errors.HasRetriableRequestsError)
579 if should_retry:
580 logger.info("Will retry failed requests.", exc_info=True)
581 logger.info("%s", exc)
582 return should_retry
583
584 try:
585 self._pending_requests = self._requests.copy()
Fang Dengf24be082018-02-10 10:09:55 -0800586 Retry(
cylan0d77ae12018-05-18 08:36:48 +0000587 _ShouldRetryHandler,
588 max_retries=self._max_retry,
Fang Dengf24be082018-02-10 10:09:55 -0800589 functor=self._ExecuteOnce,
590 sleep_multiplier=self._sleep,
591 retry_backoff_factor=self._backoff_factor)
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700592 except errors.HasRetriableRequestsError:
593 logger.debug("Some requests did not succeed after retry.")
594
595 def GetResults(self):
596 """Returns final results.
597
598 Returns:
599 results, a dictionary in the following format
600 {request_id: (response, exception)}
601 request_ids are those from requests; response
602 is the http response for the request or None on error;
603 exception is an instance of DriverError or None if no error.
604 """
605 return self._final_results
cylan31fc5332018-09-17 22:12:08 +0800606
607
608class TimeExecute(object):
609 """Count the function execute time."""
610
611 def __init__(self, function_description=None, print_before_call=True, print_status=True):
612 """Initializes the class.
613
614 Args:
615 function_description: String that describes function (e.g."Creating
616 Instance...")
617 print_before_call: Boolean, print the function description before
618 calling the function, default True.
619 print_status: Boolean, print the status of the function after the
620 function has completed, default True ("OK" or "Fail").
621 """
622 self._function_description = function_description
623 self._print_before_call = print_before_call
624 self._print_status = print_status
625
626 def __call__(self, func):
627 def DecoratorFunction(*args, **kargs):
628 """Decorator function.
629
630 Args:
631 *args: Arguments to pass to the functor.
632 **kwargs: Key-val based arguments to pass to the functor.
633
634 Raises:
635 Exception: The exception that functor(*args, **kwargs) throws.
636 """
637 timestart = time.time()
638 if self._print_before_call:
639 PrintColorString("%s ..."% self._function_description, end="")
640 try:
641 result = func(*args, **kargs)
642 if not self._print_before_call:
643 PrintColorString("%s (%ds)" % (self._function_description,
644 time.time()-timestart),
645 TextColors.OKGREEN)
646 if self._print_status:
647 PrintColorString("OK! (%ds)" % (time.time()-timestart),
648 TextColors.OKGREEN)
649 return result
650 except:
651 if self._print_status:
652 PrintColorString("Fail! (%ds)" % (time.time()-timestart),
653 TextColors.FAIL)
654 raise
655 return DecoratorFunction
cylan66713722018-10-06 01:38:26 +0800656
657
658def PickFreePort():
659 """Helper to pick a free port.
660
661 Returns:
662 Integer, a free port number.
663 """
664 tcp_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
665 tcp_socket.bind(("", 0))
666 port = tcp_socket.getsockname()[1]
667 tcp_socket.close()
668 return port
669
670
671def _ExecuteCommand(cmd, args):
672 """Execute command.
673
674 Args:
675 cmd: Strings of execute binary name.
676 args: List of args to pass in with cmd.
677
678 Raises:
679 errors.NoExecuteBin: Can't find the execute bin file.
680 """
681 bin_path = find_executable(cmd)
682 if not bin_path:
683 raise root_errors.NoExecuteCmd("unable to locate %s" % cmd)
684 command = [bin_path] + args
685 logger.debug("Running '%s'", ' '.join(command))
686 with open(os.devnull, "w") as dev_null:
687 subprocess.check_call(command, stderr=dev_null, stdout=dev_null)
688
689
690# pylint: disable=too-many-locals
691def AutoConnect(ip_addr, rsa_key_file, target_vnc_port, target_adb_port, ssh_user):
692 """Autoconnect to an AVD instance.
693
694 Args:
695 ip_addr: String, use to build the adb & vnc tunnel between local
696 and remote instance.
697 rsa_key_file: String, Private key file path to use when creating
698 the ssh tunnels.
699 target_vnc_port: Integer of target vnc port number.
700 target_adb_port: Integer of target adb port number.
701 ssh_user: String of user login into the instance.
702
703 Returns:
704 NamedTuple of (vnc_port, adb_port) SSHTUNNEL of the connect, both are
705 integers.
706 """
707 local_free_vnc_port = PickFreePort()
708 local_free_adb_port = PickFreePort()
709 try:
710 ssh_tunnel_args = _SSH_TUNNEL_ARGS % {
711 "rsa_key_file": rsa_key_file,
712 "vnc_port": local_free_vnc_port,
713 "adb_port": local_free_adb_port,
714 "target_vnc_port": target_vnc_port,
715 "target_adb_port": target_adb_port,
716 "ssh_user": ssh_user,
717 "ip_addr": ip_addr}
Kevin Cheng835a4152018-10-11 10:46:57 -0700718 _ExecuteCommand(constants.SSH_BIN, ssh_tunnel_args.split())
cylan66713722018-10-06 01:38:26 +0800719 except subprocess.CalledProcessError:
720 PrintColorString("Failed to create ssh tunnels, retry with '#acloud "
721 "reconnect'.", TextColors.FAIL)
722 try:
723 adb_connect_args = _ADB_CONNECT_ARGS % {"adb_port": local_free_adb_port}
Kevin Cheng835a4152018-10-11 10:46:57 -0700724 _ExecuteCommand(constants.ADB_BIN, adb_connect_args.split())
cylan66713722018-10-06 01:38:26 +0800725 except subprocess.CalledProcessError:
726 PrintColorString("Failed to adb connect, retry with "
727 "'#acloud reconnect'", TextColors.FAIL)
728
729 return ForwardedPorts(vnc_port=local_free_vnc_port,
730 adb_port=local_free_adb_port)
Kevin Chengeb85e862018-10-09 15:35:13 -0700731
732
733def GetAnswerFromList(answer_list, enable_choose_all=False):
734 """Get answer from a list.
735
736 Args:
737 answer_list: list of the answers to choose from.
738
739 Return:
740 List holding the answer(s).
741 """
742 print("[0] to exit.")
743 start_index = 1
744 if enable_choose_all:
745 start_index = 2
746 print("[1] for all.")
747 for num, item in enumerate(answer_list, start_index):
748 print("[%d] %s" % (num, item))
749
750 choice = -1
751 max_choice = len(answer_list) + 1
752 while True:
753 try:
754 choice = raw_input("Enter your choice[0-%d]: " % max_choice)
755 choice = int(choice)
756 except ValueError:
757 print("'%s' is not a valid integer.", choice)
758 continue
759 # Filter out choices
760 if choice == 0:
761 print("Exiting acloud.")
762 sys.exit()
763 if enable_choose_all and choice == 1:
764 return answer_list
765 if choice < 0 or choice > max_choice:
766 print("please choose between 0 and %d" % max_choice)
767 else:
768 return [answer_list[choice-start_index]]
Kevin Chengae7a49d2018-10-18 14:11:22 -0700769
770
Sam Chiu7a477f52018-10-22 11:20:36 +0800771def LaunchVNCFromReport(report, avd_spec):
772 """Launch vnc client according to the instances report.
773
774 Args:
775 report: Report object, that stores and generates report.
776 avd_spec: AVDSpec object that tells us what we're going to create.
777 """
778 for device in report.data.get("devices", []):
779 _LaunchVncClient(device.get(constants.VNC_PORT),
780 avd_width=avd_spec.hw_property["x_res"],
781 avd_height=avd_spec.hw_property["y_res"])
782
783
784def _LaunchVncClient(port=constants.DEFAULT_VNC_PORT, avd_width=None,
785 avd_height=None):
Kevin Chengae7a49d2018-10-18 14:11:22 -0700786 """Launch ssvnc.
787
788 Args:
789 port: Integer, port number.
Sam Chiu7a477f52018-10-22 11:20:36 +0800790 avd_width: String, the width of avd.
791 avd_height: String, the height of avd.
Kevin Chengae7a49d2018-10-18 14:11:22 -0700792 """
793 try:
794 os.environ[_ENV_DISPLAY]
795 except KeyError:
796 PrintColorString("Remote terminal can't support VNC. "
797 "Skipping VNC startup.", TextColors.FAIL)
798 return
799
800 if not find_executable(_VNC_BIN):
801 if GetUserAnswerYes(_CONFIRM_CONTINUE):
802 try:
803 PrintColorString("Installing ssvnc vnc client... ", end="")
804 sys.stdout.flush()
805 subprocess.check_output(_CMD_INSTALL_SSVNC, shell=True)
806 PrintColorString("Done", TextColors.OKGREEN)
807 except subprocess.CalledProcessError as cpe:
808 PrintColorString("Failed to install ssvnc: %s" %
809 cpe.output, TextColors.FAIL)
810 return
811 else:
812 return
813 ssvnc_env = os.environ.copy()
814 ssvnc_env.update(_SSVNC_ENV_VARS)
Sam Chiu7a477f52018-10-22 11:20:36 +0800815 # Override SSVNC_SCALE
816 if avd_width or avd_height:
817 scale_ratio = CalculateVNCScreenRatio(avd_width, avd_height)
818 ssvnc_env["SSVNC_SCALE"] = str(scale_ratio)
819 logger.debug("SSVNC_SCALE:%s", scale_ratio)
820
821 ssvnc_args = _CMD_START_VNC % {"bin": find_executable(_VNC_BIN),
822 "port": port}
Kevin Chengae7a49d2018-10-18 14:11:22 -0700823 subprocess.Popen(ssvnc_args.split(), env=ssvnc_env)
824
825
826def PrintDeviceSummary(report):
827 """Display summary of devices created.
828
829 -Display created device details from the report instance.
830 report example:
831 'data': [{'devices':[{'instance_name': 'ins-f6a397-none-53363',
832 'ip': u'35.234.10.162'}]}]
833 -Display error message from report.error.
834
835 Args:
836 report: A Report instance.
837 """
838 PrintColorString("\n")
839 PrintColorString("Device(s) created:")
840 for device in report.data.get("devices", []):
841 adb_serial = "(None)"
842 adb_port = device.get("adb_port")
843 if adb_port:
844 adb_serial = constants.LOCALHOST_ADB_SERIAL % adb_port
845 instance_name = device.get("instance_name")
846 instance_ip = device.get("ip")
847 instance_details = "" if not instance_name else "(%s[%s])" % (
848 instance_name, instance_ip)
849 PrintColorString(" - device serial: %s %s" % (adb_serial,
850 instance_details))
851
852 # TODO(b/117245508): Help user to delete instance if it got created.
853 if report.errors:
854 error_msg = "\n".join(report.errors)
855 PrintColorString("Fail in:\n%s\n" % error_msg, TextColors.FAIL)
Sam Chiu7a477f52018-10-22 11:20:36 +0800856
857
858def CalculateVNCScreenRatio(avd_width, avd_height):
859 """calculate the vnc screen scale ratio to fit into user's monitor.
860
861 Args:
862 avd_width: String, the width of avd.
863 avd_height: String, the height of avd.
864 Return:
865 Float, scale ratio for vnc client.
866 """
Kevin Cheng53aa5a52018-12-03 01:33:55 -0800867 try:
868 import Tkinter
869 # Some python interpreters may not be configured for Tk, just return default scale ratio.
870 except ImportError:
871 return _DEFAULT_DISPLAY_SCALE
Sam Chiu7a477f52018-10-22 11:20:36 +0800872 root = Tkinter.Tk()
873 margin = 100 # leave some space on user's monitor.
874 screen_height = root.winfo_screenheight() - margin
875 screen_width = root.winfo_screenwidth() - margin
876
Kevin Cheng53aa5a52018-12-03 01:33:55 -0800877 scale_h = _DEFAULT_DISPLAY_SCALE
878 scale_w = _DEFAULT_DISPLAY_SCALE
Sam Chiu7a477f52018-10-22 11:20:36 +0800879 if float(screen_height) < float(avd_height):
880 scale_h = round(float(screen_height) / float(avd_height), 1)
881
882 if float(screen_width) < float(avd_width):
883 scale_w = round(float(screen_width) / float(avd_width), 1)
884
885 logger.debug("scale_h: %s (screen_h: %s/avd_h: %s),"
886 " scale_w: %s (screen_w: %s/avd_w: %s)",
887 scale_h, screen_height, avd_height,
888 scale_w, screen_width, avd_width)
889
890 # Return the larger scale-down ratio.
891 return scale_h if scale_h < scale_w else scale_w
herbertxue07293a32018-11-05 20:40:11 +0800892
893
894def IsCommandRunning(command):
895 """Check if command is running.
896
897 Args:
898 command: String of command name.
899
900 Returns:
901 Boolean, True if command is running. False otherwise.
902 """
903 try:
904 with open(os.devnull, "w") as dev_null:
905 subprocess.check_call([_CMD_PGREP, command],
906 stderr=dev_null, stdout=dev_null)
907 return True
908 except subprocess.CalledProcessError:
909 return False
910
911
912def AddUserGroupsToCmd(cmd, user_groups):
913 """Add the user groups to the command if necessary.
914
915 As part of local host setup to enable local instance support, the user is
916 added to certain groups. For those settings to take effect systemwide
917 requires the user to log out and log back in. In the scenario where the
918 user has run setup and hasn't logged out, we still want them to be able to
919 launch a local instance so add the user to the groups as part of the
920 command to ensure success.
921
922 The reason using here-doc instead of '&' is all operations need to be ran in
923 ths same pid. Here's an example cmd:
924 $ sg kvm << EOF
925 sg libvirt
926 sg cvdnetwork
927 launch_cvd --cpus 2 --x_res 1280 --y_res 720 --dpi 160 --memory_mb 4096
928 EOF
929
930 Args:
931 cmd: String of the command to prepend the user groups to.
932 user_groups: List of user groups name.(String)
933
934 Returns:
935 String of the command with the user groups prepended to it if necessary,
936 otherwise the same existing command.
937 """
938 user_group_cmd = ""
939 if not CheckUserInGroups(user_groups):
940 logger.debug("Need to add user groups to the command")
941 for idx, group in enumerate(user_groups):
942 user_group_cmd += _CMD_SG + group
943 if idx == 0:
944 user_group_cmd += " <<EOF\n"
945 else:
946 user_group_cmd += "\n"
947 cmd += "\nEOF"
948 user_group_cmd += cmd
949 logger.debug("user group cmd: %s", user_group_cmd)
950 return user_group_cmd
951
952
953def CheckUserInGroups(group_name_list):
954 """Check if the current user is in the group.
955
956 Args:
957 group_name_list: The list of group name.
958 Returns:
959 True if current user is in all the groups.
960 """
961 logger.info("Checking if user is in following groups: %s", group_name_list)
962 current_groups = [grp.getgrgid(g).gr_name for g in os.getgroups()]
963 all_groups_present = True
964 for group in group_name_list:
965 if group not in current_groups:
966 all_groups_present = False
967 logger.info("missing group: %s", group)
968 return all_groups_present