def VerifyRsaPubKey(rsa):
"""Verify the format of rsa public key.
Args:
rsa: content of rsa public key. It should follow the format of
ssh-rsa AAAAB3NzaC1yc2EA.... [email protected]
Raises:
DriverError if the format is not correct.
"""
if not rsa or not all(ord(c) < 128 for c in rsa):
raise errors.DriverError(
"rsa key is empty or contains non-ascii character: %s" % rsa)
elements = rsa.split()
if len(elements) != 3:
raise errors.DriverError("rsa key is invalid, wrong format: %s" % rsa)
key_type, data, _ = elements
try:
binary_data = base64.decodestring(data)
# number of bytes of int type
int_length = 4
# binary_data is like "7ssh-key..." in a binary format.
# The first 4 bytes should represent 7, which should be
# the length of the following string "ssh-key".
# And the next 7 bytes should be string "ssh-key".
# We will verify that the rsa conforms to this format.
# ">I" in the following line means "big-endian unsigned integer".
type_length = struct.unpack(">I", binary_data[:int_length])[0]
if binary_data[int_length:int_length + type_length] != key_type:
raise errors.DriverError("rsa key is invalid: %s" % rsa)
except (struct.error, binascii.Error) as e:
raise errors.DriverError("rsa key is invalid: %s, error: %s" %
(rsa, str(e)))
def CreateSshKeyPairIfNotExist(private_key_path, public_key_path):
"""Create the ssh key pair if they don't exist.
Check if the public and private key pairs exist at
the given places. If not, create them.
Args:
private_key_path: Path to the private key file.
e.g. ~/.ssh/acloud_rsa
public_key_path: Path to the public key file.
e.g. ~/.ssh/acloud_rsa.pub
Raises:
error.DriverError: If failed to create the key pair.
"""
public_key_path = os.path.expanduser(public_key_path)
private_key_path = os.path.expanduser(private_key_path)
create_key = (not os.path.exists(public_key_path)
and not os.path.exists(private_key_path))
if not create_key:
logger.debug(
"The ssh private key (%s) or public key (%s) already exist,"
"will not automatically create the key pairs.", private_key_path,
public_key_path)
return
cmd = SSH_KEYGEN_CMD + ["-C", getpass.getuser(), "-f", private_key_path]
logger.info(
"The ssh private key (%s) and public key (%s) do not exist, "
"automatically creating key pair, calling: %s", private_key_path,
public_key_path, " ".join(cmd))
try:
subprocess.check_call(cmd, stdout=sys.stderr, stderr=sys.stdout)
except subprocess.CalledProcessError as e:
raise errors.DriverError("Failed to create ssh key pair: %s" % str(e))
except OSError as e:
raise errors.DriverError(
"Failed to create ssh key pair, please make sure "
"'ssh-keygen' is installed: %s" % str(e))
# By default ssh-keygen will create a public key file
# by append .pub to the private key file name. Rename it
# to what's requested by public_key_path.
default_pub_key_path = "%s.pub" % private_key_path
try:
if default_pub_key_path != public_key_path:
os.rename(default_pub_key_path, public_key_path)
except OSError as e:
raise errors.DriverError(
"Failed to rename %s to %s: %s" %
(default_pub_key_path, public_key_path, str(e)))
logger.info("Created ssh private key (%s) and public key (%s)",
private_key_path, public_key_path)
def CompareMachineSize(self, machine_type_1, machine_type_2, zone):
"""Compare the size of two machine types.
Args:
machine_type_1: A string representing a machine type, e.g. n1-standard-1
machine_type_2: A string representing a machine type, e.g. n1-standard-1
zone: A string representing a zone, e.g. "us-central1-f"
Returns:
1 if size of the first type is greater than the second type.
2 if size of the first type is smaller than the second type.
0 if they are equal.
Raises:
errors.DriverError: For malformed response.
"""
machine_info_1 = self.GetMachineType(machine_type_1, zone)
machine_info_2 = self.GetMachineType(machine_type_2, zone)
for metric in self.MACHINE_SIZE_METRICS:
if metric not in machine_info_1 or metric not in machine_info_2:
raise errors.DriverError(
"Malformed machine size record: Can't find '%s' in %s or %s"
% (metric, machine_info_1, machine_info_2))
if machine_info_1[metric] - machine_info_2[metric] > 0:
return 1
elif machine_info_1[metric] - machine_info_2[metric] < 0:
return -1
return 0
def AddSshRsa(self, user, ssh_rsa_path):
"""Add the public rsa key to the project's metadata.
Compute engine instances that are created after will
by default contain the key.
Args:
user: the name of the user which the key belongs to.
ssh_rsa_path: The absolute path to public rsa key.
"""
if not os.path.exists(ssh_rsa_path):
raise errors.DriverError("RSA file %s does not exist." %
ssh_rsa_path)
logger.info("Adding ssh rsa key from %s to project %s for user: %s",
ssh_rsa_path, self._project, user)
project = self.GetProject()
with open(ssh_rsa_path) as f:
rsa = f.read()
rsa = rsa.strip() if rsa else rsa
utils.VerifyRsaPubKey(rsa)
metadata = project["commonInstanceMetadata"]
for item in metadata.setdefault("items", []):
if item["key"] == "sshKeys":
sshkey_item = item
break
else:
sshkey_item = {"key": "sshKeys", "value": ""}
metadata["items"].append(sshkey_item)
entry = "%s:%s" % (user, rsa)
logger.debug("New RSA entry: %s", entry)
sshkey_item["value"] = "\n".join([sshkey_item["value"].strip(),
entry]).strip()
self.SetCommonInstanceMetadata(metadata)
def Upload(self, local_src, bucket_name, object_name, mime_type):
"""Uploads a file.
Args:
local_src: string, a local path to a file to be uploaded.
bucket_name: string, google cloud storage bucket name.
object_name: string, the name of the remote file in storage.
mime_type: string, mime-type of the file.
Returns:
URL to the inserted artifact in storage.
"""
logger.info("Uploading file: src: %s, bucket: %s, object: %s",
local_src, bucket_name, object_name)
try:
with io.FileIO(local_src, mode="rb") as fh:
media = apiclient.http.MediaIoBaseUpload(fh, mime_type)
request = self.service.objects().insert(bucket=bucket_name,
name=object_name,
media_body=media)
response = self.Execute(request)
logger.info("Uploaded artifact: %s", response["selfLink"])
return response
except OSError as e:
logger.error("Uploading artifact fails: %s", str(e))
raise errors.DriverError(str(e))
def _CreateSshKeyPairIfNecessary(cfg):
"""Create ssh key pair if necessary.
Args:
cfg: An Acloudconfig instance.
Raises:
error.DriverError: If it falls into an unexpected condition.
"""
if not cfg.ssh_public_key_path:
logger.warning("ssh_public_key_path is not specified in acloud config. "
"Project-wide public key will "
"be used when creating AVD instances. "
"Please ensure you have the correct private half of "
"a project-wide public key if you want to ssh into the "
"instances after creation.")
elif cfg.ssh_public_key_path and not cfg.ssh_private_key_path:
logger.warning("Only ssh_public_key_path is specified in acloud config,"
" but ssh_private_key_path is missing. "
"Please ensure you have the correct private half "
"if you want to ssh into the instances after creation.")
elif cfg.ssh_public_key_path and cfg.ssh_private_key_path:
utils.CreateSshKeyPairIfNotExist(
cfg.ssh_private_key_path, cfg.ssh_public_key_path)
else:
# Should never reach here.
raise errors.DriverError(
"Unexpected error in _CreateSshKeyPairIfNecessary")
def CreateDisk(self, disk_name, source_image, size_gb):
"""Create a gce disk.
Args:
disk_name: A string.
source_image: A string, name to the image name.
size_gb: Integer, size in gigabytes.
"""
if self.CheckDiskExists(disk_name, self._zone):
raise errors.DriverError(
"Failed to create disk %s, already exists." % disk_name)
if source_image and not self.CheckImageExists(source_image):
raise errors.DriverError(
"Failed to create disk %s, source image %s does not exist." %
(disk_name, source_image))
super(AndroidComputeClient, self).CreateDisk(disk_name,
source_image=source_image,
size_gb=size_gb,
zone=self._zone)
def _CreateGceImageWithLocalFile(self, local_disk_image):
"""Create a Gce image with a local image file.
The local disk image can be either a tar.gz file or a
raw vmlinux image.
e.g. /tmp/avd-system.tar.gz or /tmp/android_system_disk_syslinux.img
If a raw vmlinux image is provided, it will be archived into a tar.gz file.
The final tar.gz file will be uploaded to a cache bucket in storage.
Args:
local_disk_image: string, path to a local disk image,
Returns:
String, name of the Gce image that has been created.
Raises:
DriverError: if a file with an unexpected extension is given.
"""
logger.info("Creating a new gce image from a local file %s",
local_disk_image)
with utils.TempDir() as tempdir:
if local_disk_image.endswith(self._cfg.disk_raw_image_extension):
dest_tar_file = os.path.join(tempdir,
self._cfg.disk_image_name)
utils.MakeTarFile(
src_dict={local_disk_image: self._cfg.disk_raw_image_name},
dest=dest_tar_file)
local_disk_image = dest_tar_file
elif not local_disk_image.endswith(self._cfg.disk_image_extension):
raise errors.DriverError(
"Wrong local_disk_image type, must be a *%s file or *%s file"
% (self._cfg.disk_raw_image_extension,
self._cfg.disk_image_extension))
disk_image_id = utils.GenerateUniqueName(
suffix=self._cfg.disk_image_name)
self._storage_client.Upload(
local_src=local_disk_image,
bucket_name=self._cfg.storage_bucket_name,
object_name=disk_image_id,
mime_type=self._cfg.disk_image_mime_type)
disk_image_url = self._storage_client.GetUrl(
self._cfg.storage_bucket_name, disk_image_id)
try:
image_name = self._compute_client.GenerateImageName()
self._compute_client.CreateImage(image_name=image_name,
source_uri=disk_image_url)
finally:
self._storage_client.Delete(self._cfg.storage_bucket_name,
disk_image_id)
return image_name
def _CheckMachineSize(self):
"""Check machine size.
Check if the desired machine type |self._machine_type| meets
the requirement of minimum machine size specified as
|self._min_machine_size|.
Raises:
errors.DriverError: if check fails.
"""
if self.CompareMachineSize(self._machine_type, self._min_machine_size,
self._zone) < 0:
raise errors.DriverError(
"%s does not meet the minimum required machine size %s" %
(self._machine_type, self._min_machine_size))
def GetSerialPortOutput(self, instance, zone, port=1):
"""Get serial port output.
Args:
instance: string, instance name.
zone: string, zone name.
port: int, which COM port to read from, 1-4, default to 1.
Returns:
String, contents of the output.
Raises:
errors.DriverError: For malformed response.
"""
api = self.service.instances().getSerialPortOutput(
project=self._project, zone=zone, instance=instance, port=port)
result = self.Execute(api)
if "contents" not in result:
raise errors.DriverError(
"Malformed response for GetSerialPortOutput: %s" % result)
return result["contents"]
def _LoadSshPublicKey(ssh_public_key_path):
"""Load the content of ssh public key from a file.
Args:
ssh_public_key_path: String, path to the public key file.
E.g. ~/.ssh/acloud_rsa.pub
Returns:
String, content of the file.
Raises:
errors.DriverError if the public key file does not exist
or the content is not valid.
"""
key_path = os.path.expanduser(ssh_public_key_path)
if not os.path.exists(key_path):
raise errors.DriverError("SSH public key file %s does not exist." %
key_path)
with open(key_path) as f:
rsa = f.read()
rsa = rsa.strip() if rsa else rsa
utils.VerifyRsaPubKey(rsa)
return rsa
def _GetOperationStatus(self, operation, operation_scope, scope_name=None):
"""Get status of an operation.
Args:
operation: An Operation resource in the format of json.
operation_scope: A value from OperationScope, "zone", "region",
or "global".
scope_name: If operation_scope is "zone" or "region", this should be
the name of the zone or region, e.g. "us-central1-f".
Returns:
Status of the operation, one of "DONE", "PENDING", "RUNNING".
Raises:
errors.DriverError: if the operation fails.
"""
operation_name = operation["name"]
if operation_scope == OperationScope.GLOBAL:
api = self.service.globalOperations().get(project=self._project,
operation=operation_name)
result = self.Execute(api)
elif operation_scope == OperationScope.ZONE:
api = self.service.zoneOperations().get(project=self._project,
operation=operation_name,
zone=scope_name)
result = self.Execute(api)
elif operation_scope == OperationScope.REGION:
api = self.service.regionOperations().get(project=self._project,
operation=operation_name,
region=scope_name)
result = self.Execute(api)
if result.get("error"):
errors_list = result["error"]["errors"]
raise errors.DriverError("Get operation state failed, errors: %s" %
str(errors_list))
return result["status"]
def testCreateImageFail(self):
"""Test CreateImage fails."""
self.Patch(
gcompute_client.ComputeClient,
"WaitOnOperation",
side_effect=errors.DriverError("Expected fake error"))
self.Patch(
gcompute_client.ComputeClient,
"CheckImageExists",
return_value=True)
self.Patch(gcompute_client.ComputeClient, "DeleteImage")
resource_mock = mock.MagicMock()
self.compute_client._service.images = mock.MagicMock(
return_value=resource_mock)
resource_mock.insert = mock.MagicMock()
expected_body = {
"name": self.IMAGE,
"rawDisk": {
"source": self.GS_IMAGE_SOURCE_URI,
},
}
self.assertRaisesRegexp(
errors.DriverError,
"Expected fake error",
self.compute_client.CreateImage,
image_name=self.IMAGE,
source_uri=self.GS_IMAGE_SOURCE_URI)
resource_mock.insert.assert_called_with(
project=self.PROJECT, body=expected_body)
self.compute_client.WaitOnOperation.assert_called_with(
operation=mock.ANY,
operation_scope=gcompute_client.OperationScope.GLOBAL)
self.compute_client.CheckImageExists.assert_called_with(self.IMAGE)
self.compute_client.DeleteImage.assert_called_with(self.IMAGE)
def DownloadArtifact(self,
build_target,
build_id,
resource_id,
local_dest,
attempt_id=None):
"""Get Android build attempt information.
Args:
build_target: Target name, e.g. "gce_x86-userdebug"
build_id: Build id, a string, e.g. "2263051", "P2804227"
resource_id: Id of the resource, e.g "avd-system.tar.gz".
local_dest: A local path where the artifact should be stored.
e.g. "/tmp/avd-system.tar.gz"
attempt_id: String, attempt id, will default to DEFAULT_ATTEMPT_ID.
"""
attempt_id = attempt_id or self.DEFAULT_ATTEMPT_ID
api = self.service.buildartifact().get_media(buildId=build_id,
target=build_target,
attemptId=attempt_id,
resourceId=resource_id)
logger.info(
"Downloading artifact: target: %s, build_id: %s, "
"resource_id: %s, dest: %s", build_target, build_id, resource_id,
local_dest)
try:
with io.FileIO(local_dest, mode="wb") as fh:
downloader = apiclient.http.MediaIoBaseDownload(
fh, api, chunksize=self.DEFAULT_CHUNK_SIZE)
done = False
while not done:
_, done = downloader.next_chunk()
logger.info("Downloaded artifact: %s", local_dest)
except OSError as e:
logger.error("Downloading artifact failed: %s", str(e))
raise errors.DriverError(str(e))
def CreateDevices(self,
num,
build_target=None,
build_id=None,
gce_image=None,
local_disk_image=None,
cleanup=True,
extra_data_disk_size_gb=None,
precreated_data_image=None):
"""Creates |num| devices for given build_target and build_id.
- If gce_image is provided, will use it to create an instance.
- If local_disk_image is provided, will upload it to a temporary
caching storage bucket which is defined by user as |storage_bucket_name|
And then create an gce image with it; and then create an instance.
- If build_target and build_id are provided, will clone the disk image
via launch control to the temporary caching storage bucket.
And then create an gce image with it; and then create an instance.
Args:
num: Number of devices to create.
build_target: Target name, e.g. "gce_x86-userdebug"
build_id: Build id, a string, e.g. "2263051", "P2804227"
gce_image: string, if given, will use this image
instead of creating a new one.
implies cleanup=False.
local_disk_image: string, path to a local disk image, e.g.
/tmp/avd-system.tar.gz
cleanup: boolean, if True clean up compute engine image after creating
the instance.
extra_data_disk_size_gb: Integer, size of extra disk, or None.
precreated_data_image: A string, the image to use for the extra disk.
Raises:
errors.DriverError: If no source is specified for image creation.
"""
if gce_image:
# GCE image is provided, we can directly move to instance creation.
logger.info("Using existing gce image %s", gce_image)
image_name = gce_image
cleanup = False
elif local_disk_image:
image_name = self._CreateGceImageWithLocalFile(local_disk_image)
elif build_target and build_id:
image_name = self._CreateGceImageWithBuildInfo(build_target,
build_id)
else:
raise errors.DriverError(
"Invalid image source, must specify one of the following: gce_image, "
"local_disk_image, or build_target and build id.")
# Create GCE instances.
try:
for _ in range(num):
instance = self._compute_client.GenerateInstanceName(
build_target, build_id)
extra_disk_name = None
if extra_data_disk_size_gb > 0:
extra_disk_name = self._compute_client.GetDataDiskName(
instance)
self._compute_client.CreateDisk(extra_disk_name,
precreated_data_image,
extra_data_disk_size_gb)
self._compute_client.CreateInstance(instance, image_name,
extra_disk_name)
ip = self._compute_client.GetInstanceIP(instance)
self.devices.append(avd.AndroidVirtualDevice(
ip=ip, instance_name=instance))
finally:
if cleanup:
self._compute_client.DeleteImage(image_name)