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