blob: a90d8da62d20c61322367ef5192f047c15d836ba [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."""
Sam Chiu99dfee32018-11-20 10:19:17 +080017# pylint: disable=too-many-lines
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
Sam Chiu99dfee32018-11-20 10:19:17 +080073_EvaluatedResult = collections.namedtuple("EvaluatedResult",
74 ["is_result_ok", "result_message"])
Fang Deng69498c32017-03-02 14:29:30 -080075
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070076class TempDir(object):
Fang Deng26e4dc12018-03-04 19:01:59 -080077 """A context manager that ceates a temporary directory.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070078
Fang Deng26e4dc12018-03-04 19:01:59 -080079 Attributes:
Sam Chiu81bdc652018-06-29 18:45:08 +080080 path: The path of the temporary directory.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070081 """
82
Fang Deng26e4dc12018-03-04 19:01:59 -080083 def __init__(self):
84 self.path = tempfile.mkdtemp()
85 os.chmod(self.path, 0o700)
86 logger.debug("Created temporary dir %s", self.path)
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070087
88 def __enter__(self):
Fang Deng26e4dc12018-03-04 19:01:59 -080089 """Enter."""
90 return self.path
Keun Soo Yimb293fdb2016-09-21 16:03:44 -070091
Fang Deng26e4dc12018-03-04 19:01:59 -080092 def __exit__(self, exc_type, exc_value, traceback):
93 """Exit.
94
95 Args:
96 exc_type: Exception type raised within the context manager.
97 None if no execption is raised.
98 exc_value: Exception instance raised within the context manager.
99 None if no execption is raised.
100 traceback: Traceback for exeception that is raised within
101 the context manager.
102 None if no execption is raised.
103 Raises:
104 EnvironmentError or OSError when failed to delete temp directory.
105 """
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700106 try:
Fang Deng26e4dc12018-03-04 19:01:59 -0800107 if self.path:
108 shutil.rmtree(self.path)
109 logger.debug("Deleted temporary dir %s", self.path)
110 except EnvironmentError as e:
111 # Ignore error if there is no exception raised
112 # within the with-clause and the EnvironementError is
113 # about problem that directory or file does not exist.
114 if not exc_type and e.errno != errno.ENOENT:
115 raise
116 except Exception as e: # pylint: disable=W0703
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700117 if exc_type:
Fang Deng26e4dc12018-03-04 19:01:59 -0800118 logger.error(
cylan0d77ae12018-05-18 08:36:48 +0000119 "Encountered error while deleting %s: %s",
120 self.path,
121 str(e),
122 exc_info=True)
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700123 else:
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700124 raise
125
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700126
cylan0d77ae12018-05-18 08:36:48 +0000127def RetryOnException(retry_checker,
128 max_retries,
129 sleep_multiplier=0,
Fang Dengf24be082018-02-10 10:09:55 -0800130 retry_backoff_factor=1):
cylan0d77ae12018-05-18 08:36:48 +0000131 """Decorater which retries the function call if |retry_checker| returns true.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700132
cylan0d77ae12018-05-18 08:36:48 +0000133 Args:
134 retry_checker: A callback function which should take an exception instance
135 and return True if functor(*args, **kwargs) should be retried
136 when such exception is raised, and return False if it should
137 not be retried.
138 max_retries: Maximum number of retries allowed.
139 sleep_multiplier: Will sleep sleep_multiplier * attempt_count seconds if
140 retry_backoff_factor is 1. Will sleep
141 sleep_multiplier * (
142 retry_backoff_factor ** (attempt_count - 1))
143 if retry_backoff_factor != 1.
144 retry_backoff_factor: See explanation of sleep_multiplier.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700145
cylan0d77ae12018-05-18 08:36:48 +0000146 Returns:
147 The function wrapper.
148 """
149
150 def _Wrapper(func):
151 def _FunctionWrapper(*args, **kwargs):
152 return Retry(retry_checker, max_retries, func, sleep_multiplier,
153 retry_backoff_factor, *args, **kwargs)
154
155 return _FunctionWrapper
156
157 return _Wrapper
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700158
159
cylan4f73c1f2018-07-19 16:40:31 +0800160def Retry(retry_checker, max_retries, functor, sleep_multiplier,
161 retry_backoff_factor, *args, **kwargs):
cylan0d77ae12018-05-18 08:36:48 +0000162 """Conditionally retry a function.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700163
cylan0d77ae12018-05-18 08:36:48 +0000164 Args:
165 retry_checker: A callback function which should take an exception instance
166 and return True if functor(*args, **kwargs) should be retried
167 when such exception is raised, and return False if it should
168 not be retried.
169 max_retries: Maximum number of retries allowed.
170 functor: The function to call, will call functor(*args, **kwargs).
171 sleep_multiplier: Will sleep sleep_multiplier * attempt_count seconds if
172 retry_backoff_factor is 1. Will sleep
173 sleep_multiplier * (
174 retry_backoff_factor ** (attempt_count - 1))
175 if retry_backoff_factor != 1.
176 retry_backoff_factor: See explanation of sleep_multiplier.
177 *args: Arguments to pass to the functor.
178 **kwargs: Key-val based arguments to pass to the functor.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700179
cylan0d77ae12018-05-18 08:36:48 +0000180 Returns:
181 The return value of the functor.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700182
cylan0d77ae12018-05-18 08:36:48 +0000183 Raises:
184 Exception: The exception that functor(*args, **kwargs) throws.
185 """
186 attempt_count = 0
187 while attempt_count <= max_retries:
188 try:
189 attempt_count += 1
190 return_value = functor(*args, **kwargs)
191 return return_value
192 except Exception as e: # pylint: disable=W0703
193 if retry_checker(e) and attempt_count <= max_retries:
194 if retry_backoff_factor != 1:
195 sleep = sleep_multiplier * (retry_backoff_factor**
196 (attempt_count - 1))
197 else:
198 sleep = sleep_multiplier * attempt_count
Kevin Chengd25feee2018-05-24 10:15:20 -0700199 time.sleep(sleep)
cylan0d77ae12018-05-18 08:36:48 +0000200 else:
201 raise
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700202
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700203
Fang Dengf24be082018-02-10 10:09:55 -0800204def RetryExceptionType(exception_types, max_retries, functor, *args, **kwargs):
cylan0d77ae12018-05-18 08:36:48 +0000205 """Retry exception if it is one of the given types.
Fang Dengf24be082018-02-10 10:09:55 -0800206
cylan0d77ae12018-05-18 08:36:48 +0000207 Args:
208 exception_types: A tuple of exception types, e.g. (ValueError, KeyError)
209 max_retries: Max number of retries allowed.
210 functor: The function to call. Will be retried if exception is raised and
211 the exception is one of the exception_types.
212 *args: Arguments to pass to Retry function.
213 **kwargs: Key-val based arguments to pass to Retry functions.
Fang Dengf24be082018-02-10 10:09:55 -0800214
cylan0d77ae12018-05-18 08:36:48 +0000215 Returns:
216 The value returned by calling functor.
217 """
218 return Retry(lambda e: isinstance(e, exception_types), max_retries,
219 functor, *args, **kwargs)
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700220
221
222def PollAndWait(func, expected_return, timeout_exception, timeout_secs,
223 sleep_interval_secs, *args, **kwargs):
224 """Call a function until the function returns expected value or times out.
225
226 Args:
227 func: Function to call.
228 expected_return: The expected return value.
229 timeout_exception: Exception to raise when it hits timeout.
230 timeout_secs: Timeout seconds.
231 If 0 or less than zero, the function will run once and
232 we will not wait on it.
233 sleep_interval_secs: Time to sleep between two attemps.
234 *args: list of args to pass to func.
235 **kwargs: dictionary of keyword based args to pass to func.
236
237 Raises:
238 timeout_exception: if the run of function times out.
239 """
240 # TODO(fdeng): Currently this method does not kill
241 # |func|, if |func| takes longer than |timeout_secs|.
242 # We can use a more robust version from chromite.
243 start = time.time()
244 while True:
245 return_value = func(*args, **kwargs)
246 if return_value == expected_return:
247 return
248 elif time.time() - start > timeout_secs:
249 raise timeout_exception
250 else:
251 if sleep_interval_secs > 0:
252 time.sleep(sleep_interval_secs)
253
254
255def GenerateUniqueName(prefix=None, suffix=None):
Sam Chiu81bdc652018-06-29 18:45:08 +0800256 """Generate a random unique name using uuid4.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700257
258 Args:
259 prefix: String, desired prefix to prepend to the generated name.
260 suffix: String, desired suffix to append to the generated name.
261
262 Returns:
263 String, a random name.
264 """
265 name = uuid.uuid4().hex
266 if prefix:
267 name = "-".join([prefix, name])
268 if suffix:
269 name = "-".join([name, suffix])
270 return name
271
272
273def MakeTarFile(src_dict, dest):
274 """Archive files in tar.gz format to a file named as |dest|.
275
276 Args:
277 src_dict: A dictionary that maps a path to be archived
278 to the corresponding name that appears in the archive.
279 dest: String, path to output file, e.g. /tmp/myfile.tar.gz
280 """
281 logger.info("Compressing %s into %s.", src_dict.keys(), dest)
282 with tarfile.open(dest, "w:gz") as tar:
283 for src, arcname in src_dict.iteritems():
284 tar.add(src, arcname=arcname)
285
286
Fang Deng69498c32017-03-02 14:29:30 -0800287def CreateSshKeyPairIfNotExist(private_key_path, public_key_path):
288 """Create the ssh key pair if they don't exist.
289
cylan4f73c1f2018-07-19 16:40:31 +0800290 Case1. If the private key doesn't exist, we will create both the public key
291 and the private key.
292 Case2. If the private key exists but public key doesn't, we will create the
293 public key by using the private key.
294 Case3. If the public key exists but the private key doesn't, we will create
295 a new private key and overwrite the public key.
Fang Deng69498c32017-03-02 14:29:30 -0800296
297 Args:
298 private_key_path: Path to the private key file.
299 e.g. ~/.ssh/acloud_rsa
300 public_key_path: Path to the public key file.
301 e.g. ~/.ssh/acloud_rsa.pub
cylan4f73c1f2018-07-19 16:40:31 +0800302
Fang Deng69498c32017-03-02 14:29:30 -0800303 Raises:
304 error.DriverError: If failed to create the key pair.
305 """
306 public_key_path = os.path.expanduser(public_key_path)
307 private_key_path = os.path.expanduser(private_key_path)
cylan4f73c1f2018-07-19 16:40:31 +0800308 public_key_exist = os.path.exists(public_key_path)
309 private_key_exist = os.path.exists(private_key_path)
310 if public_key_exist and private_key_exist:
cylan0d77ae12018-05-18 08:36:48 +0000311 logger.debug(
cylan4f73c1f2018-07-19 16:40:31 +0800312 "The ssh private key (%s) and public key (%s) already exist,"
cylan0d77ae12018-05-18 08:36:48 +0000313 "will not automatically create the key pairs.", private_key_path,
314 public_key_path)
Fang Deng69498c32017-03-02 14:29:30 -0800315 return
cylan4f73c1f2018-07-19 16:40:31 +0800316 key_folder = os.path.dirname(private_key_path)
317 if not os.path.exists(key_folder):
318 os.makedirs(key_folder)
Fang Deng69498c32017-03-02 14:29:30 -0800319 try:
cylan4f73c1f2018-07-19 16:40:31 +0800320 if private_key_exist:
321 cmd = SSH_KEYGEN_PUB_CMD + ["-f", private_key_path]
322 with open(public_key_path, 'w') as outfile:
323 stream_content = subprocess.check_output(cmd)
324 outfile.write(
325 stream_content.rstrip('\n') + " " + getpass.getuser())
326 logger.info(
327 "The ssh public key (%s) do not exist, "
328 "automatically creating public key, calling: %s",
329 public_key_path, " ".join(cmd))
330 else:
331 cmd = SSH_KEYGEN_CMD + [
332 "-C", getpass.getuser(), "-f", private_key_path
333 ]
334 logger.info(
335 "Creating public key from private key (%s) via cmd: %s",
336 private_key_path, " ".join(cmd))
337 subprocess.check_call(cmd, stdout=sys.stderr, stderr=sys.stdout)
Fang Deng69498c32017-03-02 14:29:30 -0800338 except subprocess.CalledProcessError as e:
cylan0d77ae12018-05-18 08:36:48 +0000339 raise errors.DriverError("Failed to create ssh key pair: %s" % str(e))
Fang Deng69498c32017-03-02 14:29:30 -0800340 except OSError as e:
341 raise errors.DriverError(
cylan0d77ae12018-05-18 08:36:48 +0000342 "Failed to create ssh key pair, please make sure "
343 "'ssh-keygen' is installed: %s" % str(e))
Fang Deng69498c32017-03-02 14:29:30 -0800344
345 # By default ssh-keygen will create a public key file
346 # by append .pub to the private key file name. Rename it
347 # to what's requested by public_key_path.
348 default_pub_key_path = "%s.pub" % private_key_path
349 try:
350 if default_pub_key_path != public_key_path:
351 os.rename(default_pub_key_path, public_key_path)
352 except OSError as e:
353 raise errors.DriverError(
cylan0d77ae12018-05-18 08:36:48 +0000354 "Failed to rename %s to %s: %s" % (default_pub_key_path,
355 public_key_path, str(e)))
Fang Deng69498c32017-03-02 14:29:30 -0800356
357 logger.info("Created ssh private key (%s) and public key (%s)",
358 private_key_path, public_key_path)
359
360
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700361def VerifyRsaPubKey(rsa):
362 """Verify the format of rsa public key.
363
364 Args:
365 rsa: content of rsa public key. It should follow the format of
366 ssh-rsa AAAAB3NzaC1yc2EA.... test@test.com
367
368 Raises:
369 DriverError if the format is not correct.
370 """
371 if not rsa or not all(ord(c) < 128 for c in rsa):
372 raise errors.DriverError(
373 "rsa key is empty or contains non-ascii character: %s" % rsa)
374
375 elements = rsa.split()
376 if len(elements) != 3:
377 raise errors.DriverError("rsa key is invalid, wrong format: %s" % rsa)
378
379 key_type, data, _ = elements
380 try:
381 binary_data = base64.decodestring(data)
382 # number of bytes of int type
383 int_length = 4
384 # binary_data is like "7ssh-key..." in a binary format.
385 # The first 4 bytes should represent 7, which should be
386 # the length of the following string "ssh-key".
387 # And the next 7 bytes should be string "ssh-key".
388 # We will verify that the rsa conforms to this format.
389 # ">I" in the following line means "big-endian unsigned integer".
390 type_length = struct.unpack(">I", binary_data[:int_length])[0]
391 if binary_data[int_length:int_length + type_length] != key_type:
392 raise errors.DriverError("rsa key is invalid: %s" % rsa)
393 except (struct.error, binascii.Error) as e:
cylan0d77ae12018-05-18 08:36:48 +0000394 raise errors.DriverError(
395 "rsa key is invalid: %s, error: %s" % (rsa, str(e)))
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700396
chojoycecd004bc2018-09-13 10:39:00 +0800397def Decompress(sourcefile, dest=None):
398 """Decompress .zip or .tar.gz.
399
400 Args:
401 sourcefile: A string, a source file path to decompress.
402 dest: A string, a folder path as decompress destination.
403
404 Raises:
405 errors.UnsupportedCompressionFileType: Not supported extension.
406 """
407 logger.info("Start to decompress %s!", sourcefile)
408 dest_path = dest if dest else "."
409 if sourcefile.endswith(".tar.gz"):
410 with tarfile.open(sourcefile, "r:gz") as compressor:
411 compressor.extractall(dest_path)
412 elif sourcefile.endswith(".zip"):
413 with zipfile.ZipFile(sourcefile, 'r') as compressor:
414 compressor.extractall(dest_path)
415 else:
416 raise root_errors.UnsupportedCompressionFileType(
417 "Sorry, we could only support compression file type "
418 "for zip or tar.gz.")
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700419
Sam Chiu81bdc652018-06-29 18:45:08 +0800420# pylint: disable=old-style-class,no-init
421class TextColors:
422 """A class that defines common color ANSI code."""
423
424 HEADER = "\033[95m"
425 OKBLUE = "\033[94m"
426 OKGREEN = "\033[92m"
427 WARNING = "\033[93m"
428 FAIL = "\033[91m"
429 ENDC = "\033[0m"
430 BOLD = "\033[1m"
431 UNDERLINE = "\033[4m"
432
433
herbertxuedf01c422018-09-06 19:52:52 +0800434def PrintColorString(message, colors=TextColors.OKBLUE, **kwargs):
Sam Chiu81bdc652018-06-29 18:45:08 +0800435 """A helper function to print out colored text.
436
herbertxuedf01c422018-09-06 19:52:52 +0800437 Use print function "print(message, end="")" to show message in one line.
438 Example code:
439 DisplayMessages("Creating GCE instance...", end="")
440 # Job execute 20s
441 DisplayMessages("Done! (20s)")
442 Display:
443 Creating GCE instance...
444 # After job finished, messages update as following:
445 Creating GCE instance...Done! (20s)
446
Sam Chiu81bdc652018-06-29 18:45:08 +0800447 Args:
448 message: String, the message text.
449 colors: String, color code.
herbertxuedf01c422018-09-06 19:52:52 +0800450 **kwargs: dictionary of keyword based args to pass to func.
Sam Chiu81bdc652018-06-29 18:45:08 +0800451 """
herbertxuedf01c422018-09-06 19:52:52 +0800452 print(colors + message + TextColors.ENDC, **kwargs)
453 sys.stdout.flush()
Sam Chiu81bdc652018-06-29 18:45:08 +0800454
455
456def InteractWithQuestion(question, colors=TextColors.WARNING):
457 """A helper function to define the common way to run interactive cmd.
458
459 Args:
460 question: String, the question to ask user.
461 colors: String, color code.
462
463 Returns:
464 String, input from user.
465 """
466 return str(raw_input(colors + question + TextColors.ENDC).strip())
467
herbertxuedf01c422018-09-06 19:52:52 +0800468
herbertxue34776bb2018-07-03 21:57:48 +0800469def GetUserAnswerYes(question):
470 """Ask user about acloud setup question.
471
472 Args:
473 question: String, ask question for user.
474 Ex: "Are you sure to change bucket name:[y/n]"
475
476 Returns:
477 Boolean, True if answer is "Yes", False otherwise.
478 """
479 answer = InteractWithQuestion(question)
480 return answer.lower() in constants.USER_ANSWER_YES
481
Sam Chiu81bdc652018-06-29 18:45:08 +0800482
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700483class BatchHttpRequestExecutor(object):
484 """A helper class that executes requests in batch with retry.
485
486 This executor executes http requests in a batch and retry
487 those that have failed. It iteratively updates the dictionary
488 self._final_results with latest results, which can be retrieved
489 via GetResults.
490 """
491
492 def __init__(self,
493 execute_once_functor,
494 requests,
495 retry_http_codes=None,
496 max_retry=None,
497 sleep=None,
498 backoff_factor=None,
499 other_retriable_errors=None):
500 """Initializes the executor.
501
502 Args:
503 execute_once_functor: A function that execute requests in batch once.
504 It should return a dictionary like
505 {request_id: (response, exception)}
506 requests: A dictionary where key is request id picked by caller,
507 and value is a apiclient.http.HttpRequest.
508 retry_http_codes: A list of http codes to retry.
Fang Dengf24be082018-02-10 10:09:55 -0800509 max_retry: See utils.Retry.
510 sleep: See utils.Retry.
511 backoff_factor: See utils.Retry.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700512 other_retriable_errors: A tuple of error types that should be retried
513 other than errors.HttpError.
514 """
515 self._execute_once_functor = execute_once_functor
516 self._requests = requests
517 # A dictionary that maps request id to pending request.
518 self._pending_requests = {}
519 # A dictionary that maps request id to a tuple (response, exception).
520 self._final_results = {}
521 self._retry_http_codes = retry_http_codes
522 self._max_retry = max_retry
523 self._sleep = sleep
524 self._backoff_factor = backoff_factor
525 self._other_retriable_errors = other_retriable_errors
526
527 def _ShoudRetry(self, exception):
Sam Chiu81bdc652018-06-29 18:45:08 +0800528 """Check if an exception is retriable.
529
530 Args:
531 exception: An exception instance.
532 """
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700533 if isinstance(exception, self._other_retriable_errors):
534 return True
535
cylan0d77ae12018-05-18 08:36:48 +0000536 if (isinstance(exception, errors.HttpError)
537 and exception.code in self._retry_http_codes):
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700538 return True
539 return False
540
541 def _ExecuteOnce(self):
542 """Executes pending requests and update it with failed, retriable ones.
543
544 Raises:
545 HasRetriableRequestsError: if some requests fail and are retriable.
546 """
547 results = self._execute_once_functor(self._pending_requests)
548 # Update final_results with latest results.
549 self._final_results.update(results)
550 # Clear pending_requests
551 self._pending_requests.clear()
552 for request_id, result in results.iteritems():
553 exception = result[1]
554 if exception is not None and self._ShoudRetry(exception):
555 # If this is a retriable exception, put it in pending_requests
556 self._pending_requests[request_id] = self._requests[request_id]
557 if self._pending_requests:
558 # If there is still retriable requests pending, raise an error
Fang Dengf24be082018-02-10 10:09:55 -0800559 # so that Retry will retry this function with pending_requests.
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700560 raise errors.HasRetriableRequestsError(
cylan0d77ae12018-05-18 08:36:48 +0000561 "Retriable errors: %s" %
562 [str(results[rid][1]) for rid in self._pending_requests])
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700563
564 def Execute(self):
565 """Executes the requests and retry if necessary.
566
567 Will populate self._final_results.
568 """
cylan0d77ae12018-05-18 08:36:48 +0000569
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700570 def _ShouldRetryHandler(exc):
571 """Check if |exc| is a retriable exception.
572
573 Args:
574 exc: An exception.
575
576 Returns:
577 True if exception is of type HasRetriableRequestsError; False otherwise.
578 """
579 should_retry = isinstance(exc, errors.HasRetriableRequestsError)
580 if should_retry:
581 logger.info("Will retry failed requests.", exc_info=True)
582 logger.info("%s", exc)
583 return should_retry
584
585 try:
586 self._pending_requests = self._requests.copy()
Fang Dengf24be082018-02-10 10:09:55 -0800587 Retry(
cylan0d77ae12018-05-18 08:36:48 +0000588 _ShouldRetryHandler,
589 max_retries=self._max_retry,
Fang Dengf24be082018-02-10 10:09:55 -0800590 functor=self._ExecuteOnce,
591 sleep_multiplier=self._sleep,
592 retry_backoff_factor=self._backoff_factor)
Keun Soo Yimb293fdb2016-09-21 16:03:44 -0700593 except errors.HasRetriableRequestsError:
594 logger.debug("Some requests did not succeed after retry.")
595
596 def GetResults(self):
597 """Returns final results.
598
599 Returns:
600 results, a dictionary in the following format
601 {request_id: (response, exception)}
602 request_ids are those from requests; response
603 is the http response for the request or None on error;
604 exception is an instance of DriverError or None if no error.
605 """
606 return self._final_results
cylan31fc5332018-09-17 22:12:08 +0800607
608
Sam Chiu99dfee32018-11-20 10:19:17 +0800609def DefaultEvaluator(result):
610 """Default Evaluator always return result is ok.
611
612 Args:
613 result:the return value of the target function.
614 """
615 return _EvaluatedResult(is_result_ok=True, result_message=result)
616
617
618def ReportEvaluator(report):
619 """Evalute the the acloud operation by the report.
620
621 Args:
622 report:acloud.public.report() object.
623 """
624 if report is None or report.errors:
625 return _EvaluatedResult(is_result_ok=False, result_message=report.errors)
626
627 return _EvaluatedResult(is_result_ok=True, result_message=None)
628
629
cylan31fc5332018-09-17 22:12:08 +0800630class TimeExecute(object):
631 """Count the function execute time."""
632
Sam Chiu99dfee32018-11-20 10:19:17 +0800633 def __init__(self, function_description=None, print_before_call=True,
634 print_status=True, result_evaluator=DefaultEvaluator,
635 display_waiting_dots=True):
cylan31fc5332018-09-17 22:12:08 +0800636 """Initializes the class.
637
638 Args:
639 function_description: String that describes function (e.g."Creating
640 Instance...")
641 print_before_call: Boolean, print the function description before
642 calling the function, default True.
643 print_status: Boolean, print the status of the function after the
644 function has completed, default True ("OK" or "Fail").
Sam Chiu99dfee32018-11-20 10:19:17 +0800645 result_evaluator: Func object. Pass func to evaluate result.
646 Default evaluator always report result is ok and
647 failed result will be identified only in exception
648 case.
649 display_waiting_dots: Boolean, if true print the function_description
650 followed by waiting dot.
cylan31fc5332018-09-17 22:12:08 +0800651 """
652 self._function_description = function_description
653 self._print_before_call = print_before_call
654 self._print_status = print_status
Sam Chiu99dfee32018-11-20 10:19:17 +0800655 self._result_evaluator = result_evaluator
656 self._display_waiting_dots = display_waiting_dots
cylan31fc5332018-09-17 22:12:08 +0800657
658 def __call__(self, func):
659 def DecoratorFunction(*args, **kargs):
660 """Decorator function.
661
662 Args:
663 *args: Arguments to pass to the functor.
664 **kwargs: Key-val based arguments to pass to the functor.
665
666 Raises:
667 Exception: The exception that functor(*args, **kwargs) throws.
668 """
669 timestart = time.time()
670 if self._print_before_call:
Sam Chiu99dfee32018-11-20 10:19:17 +0800671 waiting_dots = "..." if self._display_waiting_dots else ""
672 PrintColorString("%s %s"% (self._function_description,
673 waiting_dots), end="")
cylan31fc5332018-09-17 22:12:08 +0800674 try:
675 result = func(*args, **kargs)
Sam Chiu99dfee32018-11-20 10:19:17 +0800676 result_time = time.time() - timestart
cylan31fc5332018-09-17 22:12:08 +0800677 if not self._print_before_call:
678 PrintColorString("%s (%ds)" % (self._function_description,
Sam Chiu99dfee32018-11-20 10:19:17 +0800679 result_time),
cylan31fc5332018-09-17 22:12:08 +0800680 TextColors.OKGREEN)
681 if self._print_status:
Sam Chiu99dfee32018-11-20 10:19:17 +0800682 evaluated_result = self._result_evaluator(result)
683 if evaluated_result.is_result_ok:
684 PrintColorString("OK! (%ds)" % (result_time),
685 TextColors.OKGREEN)
686 else:
687 PrintColorString("Fail! (%ds)" % (result_time),
688 TextColors.FAIL)
689 PrintColorString("Error: %s" %
690 evaluated_result.result_message,
691 TextColors.FAIL)
cylan31fc5332018-09-17 22:12:08 +0800692 return result
693 except:
694 if self._print_status:
695 PrintColorString("Fail! (%ds)" % (time.time()-timestart),
696 TextColors.FAIL)
697 raise
698 return DecoratorFunction
cylan66713722018-10-06 01:38:26 +0800699
700
701def PickFreePort():
702 """Helper to pick a free port.
703
704 Returns:
705 Integer, a free port number.
706 """
707 tcp_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
708 tcp_socket.bind(("", 0))
709 port = tcp_socket.getsockname()[1]
710 tcp_socket.close()
711 return port
712
713
714def _ExecuteCommand(cmd, args):
715 """Execute command.
716
717 Args:
718 cmd: Strings of execute binary name.
719 args: List of args to pass in with cmd.
720
721 Raises:
722 errors.NoExecuteBin: Can't find the execute bin file.
723 """
724 bin_path = find_executable(cmd)
725 if not bin_path:
726 raise root_errors.NoExecuteCmd("unable to locate %s" % cmd)
727 command = [bin_path] + args
728 logger.debug("Running '%s'", ' '.join(command))
729 with open(os.devnull, "w") as dev_null:
730 subprocess.check_call(command, stderr=dev_null, stdout=dev_null)
731
732
733# pylint: disable=too-many-locals
734def AutoConnect(ip_addr, rsa_key_file, target_vnc_port, target_adb_port, ssh_user):
735 """Autoconnect to an AVD instance.
736
737 Args:
738 ip_addr: String, use to build the adb & vnc tunnel between local
739 and remote instance.
740 rsa_key_file: String, Private key file path to use when creating
741 the ssh tunnels.
742 target_vnc_port: Integer of target vnc port number.
743 target_adb_port: Integer of target adb port number.
744 ssh_user: String of user login into the instance.
745
746 Returns:
747 NamedTuple of (vnc_port, adb_port) SSHTUNNEL of the connect, both are
748 integers.
749 """
750 local_free_vnc_port = PickFreePort()
751 local_free_adb_port = PickFreePort()
752 try:
753 ssh_tunnel_args = _SSH_TUNNEL_ARGS % {
754 "rsa_key_file": rsa_key_file,
755 "vnc_port": local_free_vnc_port,
756 "adb_port": local_free_adb_port,
757 "target_vnc_port": target_vnc_port,
758 "target_adb_port": target_adb_port,
759 "ssh_user": ssh_user,
760 "ip_addr": ip_addr}
Kevin Cheng835a4152018-10-11 10:46:57 -0700761 _ExecuteCommand(constants.SSH_BIN, ssh_tunnel_args.split())
cylan66713722018-10-06 01:38:26 +0800762 except subprocess.CalledProcessError:
763 PrintColorString("Failed to create ssh tunnels, retry with '#acloud "
764 "reconnect'.", TextColors.FAIL)
765 try:
766 adb_connect_args = _ADB_CONNECT_ARGS % {"adb_port": local_free_adb_port}
Kevin Cheng835a4152018-10-11 10:46:57 -0700767 _ExecuteCommand(constants.ADB_BIN, adb_connect_args.split())
cylan66713722018-10-06 01:38:26 +0800768 except subprocess.CalledProcessError:
769 PrintColorString("Failed to adb connect, retry with "
770 "'#acloud reconnect'", TextColors.FAIL)
771
772 return ForwardedPorts(vnc_port=local_free_vnc_port,
773 adb_port=local_free_adb_port)
Kevin Chengeb85e862018-10-09 15:35:13 -0700774
775
776def GetAnswerFromList(answer_list, enable_choose_all=False):
777 """Get answer from a list.
778
779 Args:
780 answer_list: list of the answers to choose from.
781
782 Return:
783 List holding the answer(s).
784 """
785 print("[0] to exit.")
786 start_index = 1
787 if enable_choose_all:
788 start_index = 2
789 print("[1] for all.")
790 for num, item in enumerate(answer_list, start_index):
791 print("[%d] %s" % (num, item))
792
793 choice = -1
794 max_choice = len(answer_list) + 1
795 while True:
796 try:
797 choice = raw_input("Enter your choice[0-%d]: " % max_choice)
798 choice = int(choice)
799 except ValueError:
800 print("'%s' is not a valid integer.", choice)
801 continue
802 # Filter out choices
803 if choice == 0:
804 print("Exiting acloud.")
805 sys.exit()
806 if enable_choose_all and choice == 1:
807 return answer_list
808 if choice < 0 or choice > max_choice:
809 print("please choose between 0 and %d" % max_choice)
810 else:
811 return [answer_list[choice-start_index]]
Kevin Chengae7a49d2018-10-18 14:11:22 -0700812
813
Sam Chiu7a477f52018-10-22 11:20:36 +0800814def LaunchVNCFromReport(report, avd_spec):
815 """Launch vnc client according to the instances report.
816
817 Args:
818 report: Report object, that stores and generates report.
819 avd_spec: AVDSpec object that tells us what we're going to create.
820 """
821 for device in report.data.get("devices", []):
822 _LaunchVncClient(device.get(constants.VNC_PORT),
823 avd_width=avd_spec.hw_property["x_res"],
824 avd_height=avd_spec.hw_property["y_res"])
825
826
827def _LaunchVncClient(port=constants.DEFAULT_VNC_PORT, avd_width=None,
828 avd_height=None):
Kevin Chengae7a49d2018-10-18 14:11:22 -0700829 """Launch ssvnc.
830
831 Args:
832 port: Integer, port number.
Sam Chiu7a477f52018-10-22 11:20:36 +0800833 avd_width: String, the width of avd.
834 avd_height: String, the height of avd.
Kevin Chengae7a49d2018-10-18 14:11:22 -0700835 """
836 try:
837 os.environ[_ENV_DISPLAY]
838 except KeyError:
839 PrintColorString("Remote terminal can't support VNC. "
840 "Skipping VNC startup.", TextColors.FAIL)
841 return
842
843 if not find_executable(_VNC_BIN):
844 if GetUserAnswerYes(_CONFIRM_CONTINUE):
845 try:
846 PrintColorString("Installing ssvnc vnc client... ", end="")
847 sys.stdout.flush()
848 subprocess.check_output(_CMD_INSTALL_SSVNC, shell=True)
849 PrintColorString("Done", TextColors.OKGREEN)
850 except subprocess.CalledProcessError as cpe:
851 PrintColorString("Failed to install ssvnc: %s" %
852 cpe.output, TextColors.FAIL)
853 return
854 else:
855 return
856 ssvnc_env = os.environ.copy()
857 ssvnc_env.update(_SSVNC_ENV_VARS)
Sam Chiu7a477f52018-10-22 11:20:36 +0800858 # Override SSVNC_SCALE
859 if avd_width or avd_height:
860 scale_ratio = CalculateVNCScreenRatio(avd_width, avd_height)
861 ssvnc_env["SSVNC_SCALE"] = str(scale_ratio)
862 logger.debug("SSVNC_SCALE:%s", scale_ratio)
863
864 ssvnc_args = _CMD_START_VNC % {"bin": find_executable(_VNC_BIN),
865 "port": port}
Kevin Chengae7a49d2018-10-18 14:11:22 -0700866 subprocess.Popen(ssvnc_args.split(), env=ssvnc_env)
867
868
869def PrintDeviceSummary(report):
870 """Display summary of devices created.
871
872 -Display created device details from the report instance.
873 report example:
874 'data': [{'devices':[{'instance_name': 'ins-f6a397-none-53363',
875 'ip': u'35.234.10.162'}]}]
876 -Display error message from report.error.
877
878 Args:
879 report: A Report instance.
880 """
881 PrintColorString("\n")
882 PrintColorString("Device(s) created:")
883 for device in report.data.get("devices", []):
884 adb_serial = "(None)"
885 adb_port = device.get("adb_port")
886 if adb_port:
887 adb_serial = constants.LOCALHOST_ADB_SERIAL % adb_port
888 instance_name = device.get("instance_name")
889 instance_ip = device.get("ip")
890 instance_details = "" if not instance_name else "(%s[%s])" % (
891 instance_name, instance_ip)
892 PrintColorString(" - device serial: %s %s" % (adb_serial,
893 instance_details))
894
895 # TODO(b/117245508): Help user to delete instance if it got created.
896 if report.errors:
897 error_msg = "\n".join(report.errors)
898 PrintColorString("Fail in:\n%s\n" % error_msg, TextColors.FAIL)
Sam Chiu7a477f52018-10-22 11:20:36 +0800899
900
901def CalculateVNCScreenRatio(avd_width, avd_height):
902 """calculate the vnc screen scale ratio to fit into user's monitor.
903
904 Args:
905 avd_width: String, the width of avd.
906 avd_height: String, the height of avd.
907 Return:
908 Float, scale ratio for vnc client.
909 """
Kevin Cheng53aa5a52018-12-03 01:33:55 -0800910 try:
911 import Tkinter
912 # Some python interpreters may not be configured for Tk, just return default scale ratio.
913 except ImportError:
914 return _DEFAULT_DISPLAY_SCALE
Sam Chiu7a477f52018-10-22 11:20:36 +0800915 root = Tkinter.Tk()
916 margin = 100 # leave some space on user's monitor.
917 screen_height = root.winfo_screenheight() - margin
918 screen_width = root.winfo_screenwidth() - margin
919
Kevin Cheng53aa5a52018-12-03 01:33:55 -0800920 scale_h = _DEFAULT_DISPLAY_SCALE
921 scale_w = _DEFAULT_DISPLAY_SCALE
Sam Chiu7a477f52018-10-22 11:20:36 +0800922 if float(screen_height) < float(avd_height):
923 scale_h = round(float(screen_height) / float(avd_height), 1)
924
925 if float(screen_width) < float(avd_width):
926 scale_w = round(float(screen_width) / float(avd_width), 1)
927
928 logger.debug("scale_h: %s (screen_h: %s/avd_h: %s),"
929 " scale_w: %s (screen_w: %s/avd_w: %s)",
930 scale_h, screen_height, avd_height,
931 scale_w, screen_width, avd_width)
932
933 # Return the larger scale-down ratio.
934 return scale_h if scale_h < scale_w else scale_w
herbertxue07293a32018-11-05 20:40:11 +0800935
936
937def IsCommandRunning(command):
938 """Check if command is running.
939
940 Args:
941 command: String of command name.
942
943 Returns:
944 Boolean, True if command is running. False otherwise.
945 """
946 try:
947 with open(os.devnull, "w") as dev_null:
948 subprocess.check_call([_CMD_PGREP, command],
949 stderr=dev_null, stdout=dev_null)
950 return True
951 except subprocess.CalledProcessError:
952 return False
953
954
955def AddUserGroupsToCmd(cmd, user_groups):
956 """Add the user groups to the command if necessary.
957
958 As part of local host setup to enable local instance support, the user is
959 added to certain groups. For those settings to take effect systemwide
960 requires the user to log out and log back in. In the scenario where the
961 user has run setup and hasn't logged out, we still want them to be able to
962 launch a local instance so add the user to the groups as part of the
963 command to ensure success.
964
965 The reason using here-doc instead of '&' is all operations need to be ran in
966 ths same pid. Here's an example cmd:
967 $ sg kvm << EOF
968 sg libvirt
969 sg cvdnetwork
970 launch_cvd --cpus 2 --x_res 1280 --y_res 720 --dpi 160 --memory_mb 4096
971 EOF
972
973 Args:
974 cmd: String of the command to prepend the user groups to.
975 user_groups: List of user groups name.(String)
976
977 Returns:
978 String of the command with the user groups prepended to it if necessary,
979 otherwise the same existing command.
980 """
981 user_group_cmd = ""
982 if not CheckUserInGroups(user_groups):
983 logger.debug("Need to add user groups to the command")
984 for idx, group in enumerate(user_groups):
985 user_group_cmd += _CMD_SG + group
986 if idx == 0:
987 user_group_cmd += " <<EOF\n"
988 else:
989 user_group_cmd += "\n"
990 cmd += "\nEOF"
991 user_group_cmd += cmd
992 logger.debug("user group cmd: %s", user_group_cmd)
993 return user_group_cmd
994
995
996def CheckUserInGroups(group_name_list):
997 """Check if the current user is in the group.
998
999 Args:
1000 group_name_list: The list of group name.
1001 Returns:
1002 True if current user is in all the groups.
1003 """
1004 logger.info("Checking if user is in following groups: %s", group_name_list)
1005 current_groups = [grp.getgrgid(g).gr_name for g in os.getgroups()]
1006 all_groups_present = True
1007 for group in group_name_list:
1008 if group not in current_groups:
1009 all_groups_present = False
1010 logger.info("missing group: %s", group)
1011 return all_groups_present