| # Copyright 2014 The Chromium OS Authors. All rights reserved. |
| # Use of this source code is governed by a BSD-style license that can be |
| # found in the LICENSE file. |
| |
| """Library for gcloud operations.""" |
| |
| from __future__ import print_function |
| |
| import re |
| |
| from chromite.compute import compute_configs |
| from chromite.lib import cros_build_lib |
| from chromite.lib import cros_logging as logging |
| from chromite.lib import timeout_util |
| |
| |
| class GCContextException(Exception): |
| """Base exception for this module.""" |
| |
| |
| # Wrap RunCommandError exceptions to distinguish the exceptions raised |
| # by commands we invoke from those that could be thrown by others who |
| # use RunCommand. |
| class GCCommandError(GCContextException, cros_build_lib.RunCommandError): |
| """Thrown when an error happened we couldn't decode.""" |
| |
| |
| class ZoneNotSpecifiedError(GCContextException): |
| """Raised when zone is not specified for a zone-specific command.""" |
| |
| |
| class GCContext(object): |
| """A wrapper around the gcloud commandline tool. |
| |
| Currently supports only `gcloud compute`. |
| """ |
| GCLOUD_BASE_COMMAND = 'gcloud' |
| GCLOUD_COMPUTE_COMMAND = 'compute' |
| |
| def __init__(self, project, zone=None, quiet=False): |
| """Initializes GCContext. |
| |
| Args: |
| project: The Google Cloud project to use. |
| zone: The default zone to operate on when zone is not given for a |
| zone-specific command. |
| quiet: If set True, skip any user prompts and use the default value. |
| """ |
| self.project = project |
| self.zone = zone |
| self.quiet = quiet |
| |
| @classmethod |
| def _GetBaseComputeCommand(cls): |
| """Returns the base Google Compute Engine command.""" |
| return [cls.GCLOUD_BASE_COMMAND, cls.GCLOUD_COMPUTE_COMMAND] |
| |
| def DoCommand(self, cmd, **kwargs): |
| """Runs |cmd|. |
| |
| cmd: The command to run. |
| project: The project to use. Defaults to using self.project. |
| kwargs: See cros_build_lib.RunCommand. |
| """ |
| cmd = self._GetBaseComputeCommand() + cmd |
| cmd += ['--project', kwargs.pop('project', self.project)] |
| if kwargs.pop('quiet', self.quiet): |
| cmd += ['--quiet'] |
| zone = kwargs.pop('zone', None) |
| if zone: |
| cmd += ['--zone', zone] |
| |
| try: |
| return cros_build_lib.RunCommand(cmd, **kwargs) |
| except cros_build_lib.RunCommandError as e: |
| raise GCCommandError(e.msg, e.result, e.exception) |
| |
| |
| def DoZoneSpecificCommand(self, cmd, **kwargs): |
| """Runs the zone-specific |cmd|. |
| |
| cmd: The command to run. |
| zone: The zone to use. Defaults to using self.zone. |
| kwargs: See DoCommand. |
| """ |
| kwargs.setdefault('zone', self.zone) |
| if not kwargs.get('zone'): |
| # To avoid ambiguity, force user to specify a zone (or a default |
| # zone) when accessing zone-specific resources. |
| raise ZoneNotSpecifiedError() |
| return self.DoCommand(cmd, **kwargs) |
| |
| def CopyFilesToInstance(self, instance, src, dest, **kwargs): |
| """Copies files from |src| to |dest| on |instance|. |
| |
| Args: |
| instance: Name of the instance. |
| src: The source path. |
| dest: The destination path. |
| kwargs: See DoCommand. |
| """ |
| return self._CopyFiles(src, '%s:%s' % (instance, dest), **kwargs) |
| |
| def CopyFilesFromInstance(self, instance, src, dest, **kwargs): |
| """Copies files from |src| on |instance| to local |dest|. |
| |
| Args: |
| instance: Name of the instance. |
| src: The source path. |
| dest: The destination path. |
| kwargs: See DoCommand. |
| """ |
| return self._CopyFiles('%s:%s' % (instance, src), dest, **kwargs) |
| |
| def _CopyFiles(self, src, dest, ssh_key_file=None, user=None, **kwargs): |
| """Copies files from |src| to |dest|. |
| |
| Args: |
| instance: Name of the instance. |
| src: The source path (a local path or user:instance@path). |
| dest: The destination path (a local path or user:instance@path). |
| kwargs: See DoCommand. |
| """ |
| command = ['copy-files'] |
| |
| if ssh_key_file: |
| command += ['--ssh-key-file', ssh_key_file] |
| if user: |
| dest = '%s@%s' % (user, dest) |
| |
| command += [src, dest] |
| |
| return self.DoZoneSpecificCommand(command, **kwargs) |
| |
| def SSH(self, instance, user=None, cmd=None, ssh_key_file=None, **kwargs): |
| """SSH into |instance|. Run |cmd| if it is provided. |
| |
| Args: |
| instance: Name of the instance. |
| cmd: Command to run on |instance|. |
| """ |
| ssh_cmd = ['ssh'] |
| if user: |
| ssh_cmd += ['%s@%s' % (user, instance)] |
| else: |
| ssh_cmd += [instance] |
| if cmd: |
| ssh_cmd += ['--command', cmd] |
| if ssh_key_file: |
| ssh_cmd += ['--ssh-key-file', ssh_key_file] |
| return self.DoZoneSpecificCommand(ssh_cmd, **kwargs) |
| |
| def ListInstances(self, **kwargs): |
| """Lists all instances.""" |
| return self.DoCommand(['instances', 'list'], **kwargs) |
| |
| def ListDisks(self, **kwargs): |
| """Lists all disks.""" |
| return self.DoCommand(['disks', 'list'], **kwargs) |
| |
| def ListImages(self, **kwargs): |
| """Lists all instances.""" |
| return self.DoCommand(['images', 'list'], **kwargs) |
| |
| def CreateImage(self, image, source_uri=None, disk=None, **kwargs): |
| """Creates an image from |source_uri| or |disk|. |
| |
| Args: |
| image: The name of image to create. |
| source_uri: The tar.gz image file (e.g. gs://foo/bar/image.tar.gz) |
| disk: The source disk to create the image from. One and only one of |
| |source_uril| and |disk| should be set. |
| kwargs: See DoCommand. |
| """ |
| if source_uri and disk: |
| raise GCContextException('Cannot specify both source uri and disk.') |
| |
| cmd = ['images', 'create', image] |
| if disk: |
| cmd += ['--source-disk', disk] |
| zone = kwargs.get('zone', self.zone) |
| if zone: |
| # Disks are zone-specific resources. |
| cmd += ['--source-disk-zone', zone] |
| |
| if source_uri: |
| cmd += ['--source-uri', source_uri] |
| |
| return self.DoCommand(cmd, **kwargs) |
| |
| def GetInstanceIP(self, instance, **kwargs): |
| """Returns an instance's ephemeral external IP address. |
| |
| May not work with all network configurations. |
| """ |
| # It sure would be nice if there were an easier way to fetch |
| # the instance's ephemeral IP address... |
| description = self.DoZoneSpecificCommand( |
| ['instances', 'describe', instance, '--format', 'text'], |
| redirect_stdout=True, **kwargs).output |
| match = re.search('natIP: ([.0-9]+)', description) |
| if match: |
| return match.group(1) |
| else: |
| raise GCContextException('Unable to parse instance IP. ' |
| '"gcloud compute instances describe %s": %s' |
| % (instance, description)) |
| |
| def CreateInstance(self, instance, image=None, machine_type=None, |
| address=None, wait_until_sshable=True, |
| scopes=None, tags=None, **kwargs): |
| """Creates an |instance|. |
| |
| Additionally, if an image is provided, adds a custom metadata pair to |
| identify the the image used to create the instance. |
| |
| Args: |
| instance: The name of the instance to create. |
| image: The source image to create |instance| from. |
| machine_type: The machine type to use. |
| address: The external IP address to assign to |instance|. |
| wait_until_sshable: After creating |instance|, wait until |
| we can ssh into |instance|. |
| scopes: The list (or tuple) of service account scopes. |
| tags: List of tags to be assigned to the instance. |
| kwargs: See DoZoneSpecificCommand. |
| """ |
| cmd = ['instances', 'create', instance] |
| if image: |
| cmd += ['--image', image] |
| cmd += ['--metadata', |
| '%s="%s"' % (compute_configs.METADATA_IMAGE_NAME, image)] |
| if address: |
| cmd += ['--address', address] |
| if machine_type: |
| cmd += ['--machine-type', machine_type] |
| if scopes: |
| cmd += ['--scopes'] + list(scopes) |
| if tags: |
| cmd += ['--tags'] + list(tags) |
| |
| ret = self.DoZoneSpecificCommand(cmd, **kwargs) |
| if wait_until_sshable: |
| def _IsUp(): |
| try: |
| instance_ip = self.GetInstanceIP(instance, **kwargs) |
| command = ['nc', '-zv', instance_ip, '22'] |
| cros_build_lib.RunCommand(command, capture_output=True) |
| except cros_build_lib.RunCommandError: |
| return False |
| else: |
| return True |
| |
| try: |
| logging.info('Waiting for the instance to be sshable...') |
| timeout = 60 * 5 |
| timeout_util.WaitForReturnTrue(_IsUp, timeout, period=5) |
| except timeout_util.TimeoutError: |
| self.DeleteInstance(instance) |
| raise GCContextException('Timed out wating to ssh into the instance') |
| |
| return ret |
| |
| def DeleteInstance(self, instance, quiet=False, keep_disks=None, **kwargs): |
| """Deletes |instance|. |
| |
| Args: |
| instance: Name of the instance. |
| quiet: Silent delete instance without prompting the user. |
| keep_disks: Keep the type of the disk; valid types are |
| 'boot', 'data', and 'all'. |
| kwargs: See DoCommand. |
| """ |
| cmd = ['instances', 'delete', instance] |
| if quiet: |
| cmd += ['-q'] |
| if keep_disks: |
| cmd += ['--keep-disks', keep_disks] |
| |
| return self.DoZoneSpecificCommand(cmd, **kwargs) |
| |
| def DeleteImage(self, image, **kwargs): |
| """Deletes |image|. User will be prompted to confirm. |
| |
| Args: |
| image: Name of the image. |
| kwargs: See DoCommand. |
| """ |
| return self.DoCommand(['images', 'delete', image], **kwargs) |
| |
| def DeleteDisk(self, disk, **kwargs): |
| """Deletes |disk|. |
| |
| Args: |
| disk: Name of the disk. |
| kwargs: See DoCommand. |
| """ |
| cmd = ['disks', 'delete', disk] |
| return self.DoZoneSpecificCommand(cmd, **kwargs) |