Esempio n. 1
0
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)))
Esempio n. 2
0
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)
Esempio n. 3
0
    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
Esempio n. 4
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)
Esempio n. 5
0
    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))
Esempio n. 6
0
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)
Esempio n. 8
0
    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))
Esempio n. 10
0
    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"]
Esempio n. 11
0
    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
Esempio n. 12
0
    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"]
Esempio n. 13
0
    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)
Esempio n. 14
0
    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))
Esempio n. 15
0
    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)