# Copyright 2015 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""A convinient wrapper of the GCE python API.

Public methods in class GceContext raise HttpError when the underlining call to
Google API fails, or gce.Error on other failures.
"""

import logging

from chromite.third_party import httplib2
from chromite.third_party.googleapiclient.discovery import build
from chromite.third_party.googleapiclient.errors import HttpError
from chromite.third_party.googleapiclient.http import HttpRequest
from chromite.third_party.oauth2client.client import GoogleCredentials

from chromite.lib import timeout_util


class Error(Exception):
    """Base exception for this module."""


class ResourceNotFoundError(Error):
    """Exceptions raised when requested GCE resource was not found."""


class RetryOnServerErrorHttpRequest(HttpRequest):
    """A HttpRequest that will be retried on server errors automatically."""

    def __init__(self, num_retries, *args, **kwargs):
        """Constructor for RetryOnServerErrorHttpRequest."""
        self.num_retries = num_retries
        super().__init__(*args, **kwargs)

    def execute(self, http=None, num_retries=None):
        """Excutes a RetryOnServerErrorHttpRequest.

        HttpRequest.execute() has the option of automatically retrying on server
        errors, i.e., 500 status codes. Call it with a non-zero value of
        |num_retries| will cause failed requests to be retried.

        Args:
            http: The httplib2.http to send this request through.
            num_retries: Number of retries. Class default value will be used if
                omitted.

        Returns:
            A deserialized object model of the response body as determined
            by the postproc. See HttpRequest.execute().
        """
        return super().execute(
            http=http, num_retries=num_retries or self.num_retries
        )


def _GetMetadataValue(metadata, key):
    """Finds a value corresponding to a given metadata key.

    Args:
        metadata: metadata object, i.e. a dict containing 'items'
            - a list of key-value pairs.
        key: name of the key.

    Returns:
        Corresponding value or None if it was not found.
    """
    for item in metadata["items"]:
        if item["key"] == key:
            return item["value"]
    return None


def _UpdateMetadataValue(metadata, key, value):
    """Updates a single key-value pair in a metadata object.

    Args:
        metadata: metadata object, i.e. a dict containing 'items'
            - a list of key-value pairs.
        key: name of the key.
        value: new value for the key, or None if it should be removed.
    """
    items = metadata.setdefault("items", [])
    for item in items:
        if item["key"] == key:
            if value is None:
                items.remove(item)
            else:
                item["value"] = value
            return

    if value is not None:
        items.append(
            {
                "key": key,
                "value": value,
            }
        )


class GceContext(object):
    """A convenient wrapper around the GCE Python API."""

    # These constants are made public so that users can customize as they need.
    DEFAULT_TIMEOUT_SEC = 5 * 60
    INSTANCE_OPERATIONS_TIMEOUT_SEC = 10 * 60
    IMAGE_OPERATIONS_TIMEOUT_SEC = 10 * 60

    _GCE_SCOPES = (
        # CreateInstance, CreateImage
        "https://www.googleapis.com/auth/compute",
        # CreateImage
        "https://www.googleapis.com/auth/devstorage.full_control",
    )
    _DEFAULT_NETWORK = "default"
    _DEFAULT_MACHINE_TYPE = "n1-standard-8"

    # Project default service account and scopes.
    _DEFAULT_SERVICE_ACCOUNT_EMAIL = "default"
    # The list is in line with what the gcloud cli uses.
    # https://cloud.google.com/sdk/gcloud/reference/compute/instances/create
    _DEFAULT_INSTANCE_SCOPES = [
        "https://www.googleapis.com/auth/cloud.useraccounts.readonly",
        "https://www.googleapis.com/auth/devstorage.read_only",
        "https://www.googleapis.com/auth/logging.write",
    ]

    # This is made public to allow easy customization of the retry behavior.
    RETRIES = 2

    def __init__(self, project, zone, credentials, thread_safe=False):
        """Initializes GceContext.

        Args:
            project: The GCP project to create instances in.
            zone: The default zone to create instances in.
            credentials: The credentials used to call the GCE API.
            thread_safe: Whether the client is expected to be thread safe.
        """
        self.project = project
        self.zone = zone

        def _BuildRequest(http, *args, **kwargs):
            """Custom request builder."""
            return self._BuildRetriableRequest(
                self.RETRIES, http, thread_safe, credentials, *args, **kwargs
            )

        self.gce_client = build(
            "compute",
            "v1",
            credentials=credentials,
            requestBuilder=_BuildRequest,
        )

        self.region = self.GetZoneRegion(zone)

    @classmethod
    def ForServiceAccount(cls, project, zone, json_key_file):
        """Creates a GceContext using service account credentials.

        About service account:
        https://developers.google.com/api-client-library/python/auth/service-accounts

        Args:
            project: The GCP project to create images and instances in.
            zone: The default zone to create instances in.
            json_key_file: Path to the service account JSON key.

        Returns:
            GceContext.
        """
        credentials = GoogleCredentials.from_stream(
            json_key_file
        ).create_scoped(cls._GCE_SCOPES)
        return GceContext(project, zone, credentials)

    @classmethod
    def ForServiceAccountThreadSafe(cls, project, zone, json_key_file):
        """Creates a thread-safe GceContext using service account credentials.

        About service account:
        https://developers.google.com/api-client-library/python/auth/service-accounts

        Args:
            project: The GCP project to create images and instances in.
            zone: The default zone to create instances in.
            json_key_file: Path to the service account JSON key.

        Returns:
            GceContext.
        """
        credentials = GoogleCredentials.from_stream(
            json_key_file
        ).create_scoped(cls._GCE_SCOPES)
        return GceContext(project, zone, credentials, thread_safe=True)

    def CreateAddress(self, name, region=None):
        """Reserves an external IP address.

        Args:
            name: The name to assign to the address.
            region: Region to reserve the address in.

        Returns:
            The reserved address as a string.
        """
        body = {
            "name": name,
        }
        operation = (
            self.gce_client.addresses()
            .insert(
                project=self.project, region=region or self.region, body=body
            )
            .execute()
        )
        self._WaitForRegionOperation(
            operation["name"],
            region,
            timeout_sec=self.INSTANCE_OPERATIONS_TIMEOUT_SEC,
        )

        address = (
            self.gce_client.addresses()
            .get(
                project=self.project, region=region or self.region, address=name
            )
            .execute()
        )

        return address["address"]

    def DeleteAddress(self, name, region=None):
        """Frees up an external IP address.

        Args:
            name: The name of the address.
            region: Region of the address.
        """
        operation = (
            self.gce_client.addresses()
            .delete(
                project=self.project, region=region or self.region, address=name
            )
            .execute()
        )
        self._WaitForRegionOperation(
            operation["name"],
            region=region or self.region,
            timeout_sec=self.INSTANCE_OPERATIONS_TIMEOUT_SEC,
        )

    def GetZoneRegion(self, zone=None):
        """Resolves name of the region that a zone belongs to.

        Args:
            zone: The zone to resolve.

        Returns:
            Name of the region corresponding to the zone.
        """
        zone_resource = (
            self.gce_client.zones()
            .get(project=self.project, zone=zone or self.zone)
            .execute()
        )
        return zone_resource["region"].split("/")[-1]

    def CreateInstance(
        self,
        name,
        image,
        zone=None,
        network=None,
        subnet=None,
        machine_type=None,
        default_scopes=True,
        static_address=None,
        **kwargs,
    ):
        """Creates an instance with the given image and waits until it's ready.

        Args:
            name: Instance name.
            image: Fully spelled URL of the image, e.g., for private images,
                'global/images/my-private-image', or for images from a
                publicly-available project,
                'projects/debian-cloud/global/images/debian-7-wheezy-vYYYYMMDD'.
                Details:
                https://cloud.google.com/compute/docs/reference/latest/instances/insert
            zone: The zone to create the instance in. Default zone will be used
                if omitted.
            network: An existing network to create the instance in. Default
                network will be used if omitted.
            subnet: The subnet to create the instance in.
            machine_type: The machine type to use. Default machine type will be
                used if omitted.
            default_scopes: If true, the default scopes are added to the
                instances.
            static_address: External IP address to assign to the instance as a
                string. If None an ephemeral address will be used.
            **kwargs: Other possible Instance Resource properties.
                https://cloud.google.com/compute/docs/reference/latest/instances#resource
                Note that values from kwargs will overrule properties
                constructed from positional arguments, i.e., name, image, zone,
                network and machine_type.

        Returns:
            URL to the created instance.
        """
        logging.info('Creating instance "%s" with image "%s" ...', name, image)
        network = "global/networks/%s" % network or self._DEFAULT_NETWORK
        machine_type = "zones/%s/machineTypes/%s" % (
            zone or self.zone,
            machine_type or self._DEFAULT_MACHINE_TYPE,
        )
        service_accounts = (
            (
                {
                    "email": self._DEFAULT_SERVICE_ACCOUNT_EMAIL,
                    "scopes": self._DEFAULT_INSTANCE_SCOPES,
                },
            )
            if default_scopes
            else ()
        )

        config = {
            "name": name,
            "machineType": machine_type,
            "disks": (
                {
                    "boot": True,
                    "autoDelete": True,
                    "initializeParams": {
                        "sourceImage": image,
                    },
                },
            ),
            "networkInterfaces": (
                {
                    "network": network,
                    "accessConfigs": (
                        {
                            "type": "ONE_TO_ONE_NAT",
                            "name": "External NAT",
                        },
                    ),
                },
            ),
            "serviceAccounts": service_accounts,
        }
        config.update(**kwargs)
        if static_address is not None:
            config["networkInterfaces"][0]["accessConfigs"][0][
                "natIP"
            ] = static_address
        if subnet is not None:
            region = self.GetZoneRegion(zone)
            config["networkInterfaces"][0][
                "subnetwork"
            ] = "regions/%s/subnetworks/%s" % (region, subnet)
        operation = (
            self.gce_client.instances()
            .insert(project=self.project, zone=zone or self.zone, body=config)
            .execute()
        )
        self._WaitForZoneOperation(
            operation["name"],
            timeout_sec=self.INSTANCE_OPERATIONS_TIMEOUT_SEC,
            timeout_handler=lambda: self.DeleteInstance(name),
        )
        return operation["targetLink"]

    def DeleteInstance(self, name, zone=None):
        """Deletes an instance with the name and waits until it's done.

        Args:
            name: Name of the instance to delete.
            zone: Zone where the instance is in. Default zone will be used if
                omitted.
        """
        logging.info('Deleting instance "%s" ...', name)
        operation = (
            self.gce_client.instances()
            .delete(project=self.project, zone=zone or self.zone, instance=name)
            .execute()
        )
        self._WaitForZoneOperation(
            operation["name"], timeout_sec=self.INSTANCE_OPERATIONS_TIMEOUT_SEC
        )

    def StartInstance(self, name, zone=None):
        """Starts an instance with the name and waits until it's done.

        Args:
            name: Name of the instance to start.
            zone: Zone where the instance is in. Default zone will be used if
                omitted.
        """
        logging.info('Starting instance "%s" ...', name)
        operation = (
            self.gce_client.instances()
            .start(project=self.project, zone=zone or self.zone, instance=name)
            .execute()
        )
        self._WaitForZoneOperation(
            operation["name"], timeout_sec=self.INSTANCE_OPERATIONS_TIMEOUT_SEC
        )

    def StopInstance(self, name, zone=None):
        """Stops an instance with the name and waits until it's done.

        Args:
            name: Name of the instance to stop.
            zone: Zone where the instance is in. Default zone will be used if
                omitted.
        """
        logging.info('Stopping instance "%s" ...', name)
        operation = (
            self.gce_client.instances()
            .stop(project=self.project, zone=zone or self.zone, instance=name)
            .execute()
        )
        self._WaitForZoneOperation(
            operation["name"], timeout_sec=self.INSTANCE_OPERATIONS_TIMEOUT_SEC
        )

    def CreateImage(self, name, source):
        """Creates an image with the given |source|.

        Args:
            name: Name of the image to be created.
            source: Google Cloud Storage object of the source disk. e.g.:
                https://storage.googleapis.com/my-gcs-bucket/test_image.tar.gz

        Returns:
            URL to the created image.
        """
        logging.info('Creating image "%s" with "source" %s ...', name, source)
        config = {
            "name": name,
            "rawDisk": {
                "source": source,
            },
        }
        operation = (
            self.gce_client.images()
            .insert(project=self.project, body=config)
            .execute()
        )
        self._WaitForGlobalOperation(
            operation["name"],
            timeout_sec=self.IMAGE_OPERATIONS_TIMEOUT_SEC,
            timeout_handler=lambda: self.DeleteImage(name),
        )
        return operation["targetLink"]

    def DeleteImage(self, name):
        """Deletes an image and waits until it's deleted.

        Args:
            name: Name of the image to delete.
        """
        logging.info('Deleting image "%s" ...', name)
        operation = (
            self.gce_client.images()
            .delete(project=self.project, image=name)
            .execute()
        )
        self._WaitForGlobalOperation(
            operation["name"], timeout_sec=self.IMAGE_OPERATIONS_TIMEOUT_SEC
        )

    def ListInstances(self, zone=None):
        """Lists all instances.

        Args:
            zone: Zone where the instances are in. Default zone will be used if
                omitted.

        Returns:
            A list of Instance Resources if found, or an empty list otherwise.
        """
        result = (
            self.gce_client.instances()
            .list(project=self.project, zone=zone or self.zone)
            .execute()
        )
        return result.get("items", [])

    def ListImages(self):
        """Lists all images.

        Returns:
            A list of Image Resources if found, or an empty list otherwise.
        """
        result = self.gce_client.images().list(project=self.project).execute()
        return result.get("items", [])

    def GetInstance(self, instance, zone=None):
        """Gets an Instance Resource by name and zone.

        Args:
            instance: Name of the instance.
            zone: Zone where the instance is in. Default zone will be used if
                omitted.

        Returns:
            An Instance Resource.

        Raises:
            ResourceNotFoundError: if instance was not found.
            HttpError: other HTTP failures.
        """
        try:
            return (
                self.gce_client.instances()
                .get(
                    project=self.project,
                    zone=zone or self.zone,
                    instance=instance,
                )
                .execute()
            )
        except HttpError as e:
            if e.resp.status == 404:
                raise ResourceNotFoundError(
                    'Instance "%s" for project "%s" in zone "%s" was not found.'
                    % (instance, self.project, zone or self.zone)
                )
            else:
                raise

    def GetInstanceIP(self, instance, zone=None):
        """Gets the external IP of an instance.

        Args:
            instance: Name of the instance to get IP for.
            zone: Zone where the instance is in. Default zone will be used if
                omitted.

        Returns:
            External IP address of the instance.

        Raises:
            Error: Something went wrong when trying to get IP for the instance.
        """
        result = self.GetInstance(instance, zone)
        try:
            return result["networkInterfaces"][0]["accessConfigs"][0]["natIP"]
        except (KeyError, IndexError):
            raise Error("Failed to get IP address for instance %s" % instance)

    def GetInstanceInternalIP(self, instance, zone=None):
        """Gets the internal IP of an instance."""
        result = self.GetInstance(instance, zone)
        try:
            return result["networkInterfaces"][0]["networkIP"]
        except (KeyError, IndexError):
            raise Error("Failed to get internal IP for instance %s" % instance)

    def GetImage(self, image):
        """Gets an Image Resource by name.

        Args:
            image: Name of the image to look for.

        Returns:
            An Image Resource.

        Raises:
            ResourceNotFoundError: The requested image was not found.
        """
        try:
            return (
                self.gce_client.images()
                .get(project=self.project, image=image)
                .execute()
            )
        except HttpError as e:
            if e.resp.status == 404:
                raise ResourceNotFoundError(
                    'Image "%s" for project "%s" was not found.'
                    % (image, self.project)
                )
            else:
                raise

    def InstanceExists(self, instance, zone=None):
        """Checks if an instance exists in the current project.

        Args:
            instance: Name of the instance to check existence of.
            zone: Zone where the instance is in. Default zone will be used if
                omitted.

        Returns:
            True if the instance exists or False otherwise.
        """
        try:
            return self.GetInstance(instance, zone) is not None
        except ResourceNotFoundError:
            return False

    def ImageExists(self, image):
        """Checks if an image exists in the current project.

        Args:
            image: Name of the image to check existence of.

        Returns:
            True if the instance exists or False otherwise.
        """
        try:
            return self.GetImage(image) is not None
        except ResourceNotFoundError:
            return False

    def GetCommonInstanceMetadata(self, key):
        """Looks up a single project metadata value.

        Args:
            key: Metadata key name.

        Returns:
            Metadata value corresponding to the key, or None if it was not
            found.
        """
        projects_data = (
            self.gce_client.projects().get(project=self.project).execute()
        )
        metadata = projects_data["commonInstanceMetadata"]
        return _GetMetadataValue(metadata, key)

    def SetCommonInstanceMetadata(self, key, value):
        """Sets a single project metadata value.

        Args:
            key: Metadata key to be set.
            value: New value, or None if the given key should be removed.
        """
        projects_data = (
            self.gce_client.projects().get(project=self.project).execute()
        )
        metadata = projects_data["commonInstanceMetadata"]
        _UpdateMetadataValue(metadata, key, value)
        operation = (
            self.gce_client.projects()
            .setCommonInstanceMetadata(project=self.project, body=metadata)
            .execute()
        )
        self._WaitForGlobalOperation(operation["name"])

    def GetInstanceMetadata(self, instance, key):
        """Looks up instance's metadata value.

        Args:
            instance: Name of the instance.
            key: Metadata key name.

        Returns:
            Metadata value corresponding to the key, or None if it was not
            found.
        """
        instance_data = self.GetInstance(instance)
        metadata = instance_data["metadata"]
        return _GetMetadataValue(metadata, key)

    def SetInstanceMetadata(self, instance, key, value):
        """Sets a single instance metadata value.

        Args:
            instance: Name of the instance.
            key: Metadata key to be set.
            value: New value, or None if the given key should be removed.
        """
        instance_data = self.GetInstance(instance)
        metadata = instance_data["metadata"]
        _UpdateMetadataValue(metadata, key, value)
        operation = (
            self.gce_client.instances()
            .setMetadata(
                project=self.project,
                zone=self.zone,
                instance=instance,
                body=metadata,
            )
            .execute()
        )
        self._WaitForZoneOperation(operation["name"])

    def _WaitForZoneOperation(
        self, operation, zone=None, timeout_sec=None, timeout_handler=None
    ):
        """Waits until a GCE ZoneOperation is finished or timed out.

        Args:
            operation: The GCE operation to wait for.
            zone: The zone that |operation| belongs to.
            timeout_sec: The maximum number of seconds to wait for.
            timeout_handler: A callable to be executed when timeout happens.

        Raises:
            Error when timeout happens or the operation fails.
        """
        get_request = self.gce_client.zoneOperations().get(
            project=self.project, zone=zone or self.zone, operation=operation
        )
        self._WaitForOperation(
            operation, get_request, timeout_sec, timeout_handler=timeout_handler
        )

    def _WaitForRegionOperation(
        self, operation, region, timeout_sec=None, timeout_handler=None
    ):
        """Waits until a GCE RegionOperation is finished or timed out.

        Args:
            operation: The GCE operation to wait for.
            region: The region that |operation| belongs to.
            timeout_sec: The maximum number of seconds to wait for.
            timeout_handler: A callable to be executed when timeout happens.

        Raises:
            Error when timeout happens or the operation fails.
        """
        get_request = self.gce_client.regionOperations().get(
            project=self.project,
            region=region or self.region,
            operation=operation,
        )
        self._WaitForOperation(
            operation, get_request, timeout_sec, timeout_handler=timeout_handler
        )

    def _WaitForGlobalOperation(
        self, operation, timeout_sec=None, timeout_handler=None
    ):
        """Waits until a GCE GlobalOperation is finished or timed out.

        Args:
            operation: The GCE operation to wait for.
            timeout_sec: The maximum number of seconds to wait for.
            timeout_handler: A callable to be executed when timeout happens.

        Raises:
            Error when timeout happens or the operation fails.
        """
        get_request = self.gce_client.globalOperations().get(
            project=self.project, operation=operation
        )
        self._WaitForOperation(
            operation,
            get_request,
            timeout_sec=timeout_sec,
            timeout_handler=timeout_handler,
        )

    def _WaitForOperation(
        self,
        operation,
        get_operation_request,
        timeout_sec=None,
        timeout_handler=None,
    ):
        """Waits until timeout, or it receives a response with a 'DONE' status.

        Args:
            operation: The GCE operation to wait for.
            get_operation_request: The HTTP request to get the operation's
                status. This request will be executed periodically until it
                returns a status 'DONE'.
            timeout_sec: The maximum number of seconds to wait for.
            timeout_handler: A callable to be executed when times out.

        Raises:
            Error when timeout happens or the operation fails.
        """

        def _IsDone():
            result = get_operation_request.execute()
            if result["status"] == "DONE":
                if "error" in result:
                    raise Error(result["error"])
                return True
            return False

        try:
            timeout = timeout_sec or self.DEFAULT_TIMEOUT_SEC
            logging.info(
                "Waiting up to %d seconds for operation [%s] to complete...",
                timeout,
                operation,
            )
            timeout_util.WaitForReturnTrue(_IsDone, timeout, period=1)
        except timeout_util.TimeoutError:
            if timeout_handler:
                timeout_handler()
            raise Error(
                "Timeout waiting for operation [%s] to complete" % operation
            )

    def _BuildRetriableRequest(
        self,
        num_retries,
        http,
        thread_safe=False,
        credentials=None,
        *args,
        **kwargs,
    ):
        """Builds a request that will be automatically retried on server errors.

        Args:
            num_retries: The maximum number of times to retry until give up.
            http: An httplib2.Http object that this request will be executed
                through.
            thread_safe: Whether the request needs to be thread-safe.
            credentials: Credentials to apply to the request.
            *args: Optional positional arguments.
            **kwargs: Optional keyword arguments.

        Returns:
            RetryOnServerErrorHttpRequest: A request that will automatically
                retry on server errors.
        """
        if thread_safe:
            # Create a new http object for every request.
            http = credentials.authorize(httplib2.Http())
        return RetryOnServerErrorHttpRequest(num_retries, http, *args, **kwargs)