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(six.b(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] != six.b(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 CreateDisk(self,
disk_name,
source_image,
size_gb,
zone=None,
source_project=None,
disk_type=gcompute_client.PersistentDiskType.STANDARD):
"""Create a gce disk.
Args:
disk_name: String, name of disk.
source_image: String, name to the image name.
size_gb: Integer, size in gigabytes.
zone: String, name of the zone, e.g. us-central1-b.
source_project: String, required if the image is located in a different
project.
disk_type: String, a value from PersistentDiskType, STANDARD
for regular hard disk or SSD for solid state disk.
"""
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=zone or self._zone)
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 upload_file:
media = apiclient.http.MediaIoBaseUpload(
upload_file, 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 _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 _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 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. "aosp_cf_x86_phone-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, apiclient.errors.HttpError) 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,
avd_spec=None,
extra_scopes=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. "aosp_cf_x86_phone-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.
avd_spec: AVDSpec object for pass hw_property.
extra_scopes: A list of extra scopes given to the new instance.
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=instance,
image_name=image_name,
extra_disk_name=extra_disk_name,
avd_spec=avd_spec,
extra_scopes=extra_scopes)
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)
def CreateSshKeyPairIfNotExist(private_key_path, public_key_path):
"""Create the ssh key pair if they don't exist.
Case1. If the private key doesn't exist, we will create both the public key
and the private key.
Case2. If the private key exists but public key doesn't, we will create the
public key by using the private key.
Case3. If the public key exists but the private key doesn't, we will create
a new private key and overwrite the public key.
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)
public_key_exist = os.path.exists(public_key_path)
private_key_exist = os.path.exists(private_key_path)
if public_key_exist and private_key_exist:
logger.debug(
"The ssh private key (%s) and public key (%s) already exist,"
"will not automatically create the key pairs.", private_key_path,
public_key_path)
return
key_folder = os.path.dirname(private_key_path)
if not os.path.exists(key_folder):
os.makedirs(key_folder)
try:
if private_key_exist:
cmd = SSH_KEYGEN_PUB_CMD + ["-f", private_key_path]
with open(public_key_path, 'w') as outfile:
stream_content = CheckOutput(cmd)
outfile.write(
stream_content.rstrip('\n') + " " + getpass.getuser())
logger.info(
"The ssh public key (%s) do not exist, "
"automatically creating public key, calling: %s",
public_key_path, " ".join(cmd))
else:
cmd = SSH_KEYGEN_CMD + [
"-C", getpass.getuser(), "-f", private_key_path
]
logger.info(
"Creating public key from private key (%s) via cmd: %s",
private_key_path, " ".join(cmd))
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)