infra/bots/assets/asset_utils.py

#!/usr/bin/env python
#
# Copyright 2016 Google Inc.
#
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.


"""Utilities for managing assets."""


import argparse
import json
import os
import shlex
import shutil
import subprocess
import sys

INFRA_BOTS_DIR = os.path.abspath(os.path.realpath(os.path.join(
    os.path.dirname(os.path.abspath(__file__)), os.pardir)))
sys.path.insert(0, INFRA_BOTS_DIR)
import utils
import zip_utils


ASSETS_DIR = os.path.join(INFRA_BOTS_DIR, 'assets')
SKIA_DIR = os.path.abspath(os.path.join(INFRA_BOTS_DIR, os.pardir, os.pardir))

CIPD_PACKAGE_NAME_TMPL = 'skia/bots/%s'
DEFAULT_CIPD_SERVICE_URL = 'https://chrome-infra-packages.appspot.com'

DEFAULT_GS_BUCKET = 'skia-assets'
GS_SUBDIR_TMPL = 'gs://%s/assets/%s'
GS_PATH_TMPL = '%s/%s.zip'

TAG_PROJECT_SKIA = 'project:skia'
TAG_VERSION_PREFIX = 'version:'
TAG_VERSION_TMPL = '%s%%s' % TAG_VERSION_PREFIX

VERSION_FILENAME = 'VERSION'
ZIP_BLACKLIST = ['.git', '.svn', '*.pyc', '.DS_STORE']


class CIPDStore(object):
  """Wrapper object for CIPD."""
  def __init__(self, cipd_url=DEFAULT_CIPD_SERVICE_URL,
               service_account_json=None):
    self._cipd = 'cipd'
    if sys.platform == 'win32':
      self._cipd = 'cipd.bat'
    self._cipd_url = cipd_url
    if service_account_json:
      self._service_account_json = os.path.abspath(service_account_json)
    else:
      self._service_account_json = None
    self._check_setup()

  def _check_setup(self):
    """Verify that we have the CIPD binary and that we're authenticated."""
    try:
      self._run(['auth-info'], specify_service_url=False)
    except OSError:
      raise Exception('CIPD binary not found on your path (typically in '
                      'depot_tools). You may need to update depot_tools.')
    except subprocess.CalledProcessError:
      raise Exception('CIPD not authenticated. You may need to run:\n\n'
                      '$ %s auth-login' % self._cipd)

  def _run(self, cmd, specify_service_url=True):
    """Run the given command."""
    cipd_args = []
    if specify_service_url:
      cipd_args.extend(['--service-url', self._cipd_url])
    if self._service_account_json:
      cipd_args.extend(['-service-account-json', self._service_account_json])
    elif os.getenv('USE_CIPD_GCE_AUTH'):
      # Enable automatic GCE authentication. For context see
      # https://bugs.chromium.org/p/skia/issues/detail?id=6385#c3
      cipd_args.extend(['-service-account-json', ':gce'])
    return subprocess.check_output(
        [self._cipd] + cmd + cipd_args,
        stderr=subprocess.STDOUT)

  def _json_output(self, cmd):
    """Run the given command, return the JSON output."""
    with utils.tmp_dir():
      json_output = os.path.join(os.getcwd(), 'output.json')
      self._run(cmd + ['--json-output', json_output])
      with open(json_output) as f:
        parsed = json.load(f)
    return parsed.get('result', [])

  def _search(self, pkg_name):
    try:
      res = self._json_output(['search', pkg_name, '--tag', TAG_PROJECT_SKIA])
    except subprocess.CalledProcessError as e:
      if 'no such package' in e.output:
        return []
      raise
    return [r['instance_id'] for r in res or []]

  def _describe(self, pkg_name, instance_id):
    """Obtain details about the given package and instance ID."""
    return self._json_output(['describe', pkg_name, '--version', instance_id])

  def get_available_versions(self, name):
    """List available versions of the asset."""
    pkg_name = CIPD_PACKAGE_NAME_TMPL % name
    versions = []
    for instance_id in self._search(pkg_name):
      details = self._describe(pkg_name, instance_id)
      for tag in details.get('tags'):
        tag_name = tag.get('tag', '')
        if tag_name.startswith(TAG_VERSION_PREFIX):
          trimmed = tag_name[len(TAG_VERSION_PREFIX):]
          try:
            versions.append(int(trimmed))
          except ValueError:
            raise ValueError('Found package instance with invalid version '
                             'tag: %s' % tag_name)
    versions.sort()
    return versions

  def upload(self, name, version, target_dir, extra_tags=None):
    """Create a CIPD package."""
    cmd = [
        'create',
        '--name', CIPD_PACKAGE_NAME_TMPL % name,
        '--in', target_dir,
        '--tag', TAG_PROJECT_SKIA,
        '--tag', TAG_VERSION_TMPL % version,
        '--compression-level', '1',
        '-verification-timeout', '30m0s',
    ]
    if extra_tags:
      for tag in extra_tags:
        cmd.extend(['--tag', tag])
    self._run(cmd)

  def download(self, name, version, target_dir):
    """Download a CIPD package."""
    pkg_name = CIPD_PACKAGE_NAME_TMPL % name
    version_tag = TAG_VERSION_TMPL % version
    target_dir = os.path.abspath(target_dir)
    with utils.tmp_dir():
      infile = os.path.join(os.getcwd(), 'input')
      with open(infile, 'w') as f:
        f.write('%s %s' % (pkg_name, version_tag))
      self._run([
          'ensure',
          '--root', target_dir,
          '--list', infile,
      ])

  def delete_contents(self, name):
    """Delete data for the given asset."""
    self._run(['pkg-delete', CIPD_PACKAGE_NAME_TMPL % name])


class GSStore(object):
  """Wrapper object for interacting with Google Storage."""
  def __init__(self, gsutil=None, bucket=DEFAULT_GS_BUCKET,
               service_account_json=None):
    if gsutil:
      gsutil = os.path.abspath(gsutil)
    else:
      gsutils = subprocess.check_output([
          utils.WHICH, 'gsutil']).rstrip().splitlines()
      for g in gsutils:
        ok = True
        try:
          subprocess.check_call([g, 'version'])
        except OSError:
          ok = False
        if ok:
          gsutil = g
          break
    self._gsutil = [gsutil]
    if gsutil.endswith('.py'):
      self._gsutil = ['python', gsutil]
    if service_account_json:
      sa = os.path.abspath(service_account_json)
      self._gsutil += ['-o', 'Credentials:gs_service_key_file=' + sa]
    self._gs_bucket = bucket

  def copy(self, src, dst):
    """Copy src to dst."""
    subprocess.check_call(self._gsutil + ['cp', src, dst])

  def list(self, path):
    """List objects in the given path."""
    try:
      return subprocess.check_output(self._gsutil + ['ls', path]).splitlines()
    except subprocess.CalledProcessError:
      # If the prefix does not exist, we'll get an error, which is okay.
      return []

  def get_available_versions(self, name):
    """Return the existing version numbers for the asset."""
    files = self.list(GS_SUBDIR_TMPL % (self._gs_bucket, name))
    bnames = [os.path.basename(f) for f in files]
    suffix = '.zip'
    versions = [int(f[:-len(suffix)]) for f in bnames if f.endswith(suffix)]
    versions.sort()
    return versions

  # pylint: disable=unused-argument
  def upload(self, name, version, target_dir, extra_tags=None):
    """Upload to GS."""
    target_dir = os.path.abspath(target_dir)
    with utils.tmp_dir():
      zip_file = os.path.join(os.getcwd(), '%d.zip' % version)
      zip_utils.zip(target_dir, zip_file, blacklist=ZIP_BLACKLIST)
      gs_path = GS_PATH_TMPL % (GS_SUBDIR_TMPL % (self._gs_bucket, name),
                                str(version))
      self.copy(zip_file, gs_path)

  def download(self, name, version, target_dir):
    """Download from GS."""
    gs_path = GS_PATH_TMPL % (GS_SUBDIR_TMPL % (self._gs_bucket, name),
                              str(version))
    target_dir = os.path.abspath(target_dir)
    with utils.tmp_dir():
      zip_file = os.path.join(os.getcwd(), '%d.zip' % version)
      self.copy(gs_path, zip_file)
      zip_utils.unzip(zip_file, target_dir)

  def delete_contents(self, name):
    """Delete data for the given asset."""
    gs_path = GS_SUBDIR_TMPL % (self._gs_bucket, name)
    attempt_delete = True
    try:
      subprocess.check_call(self._gsutil + ['ls', gs_path])
    except subprocess.CalledProcessError:
      attempt_delete = False
    if attempt_delete:
      subprocess.check_call(self._gsutil + ['rm', '-rf', gs_path])


class MultiStore(object):
  """Wrapper object which uses CIPD as the primary store and GS for backup."""
  def __init__(self, cipd_url=DEFAULT_CIPD_SERVICE_URL,
               service_account_json=None,
               gsutil=None, gs_bucket=DEFAULT_GS_BUCKET):
    self._cipd = CIPDStore(cipd_url=cipd_url,
                           service_account_json=service_account_json)
    self._gs = GSStore(gsutil=gsutil, bucket=gs_bucket,
                       service_account_json=service_account_json)

  def get_available_versions(self, name):
    return self._cipd.get_available_versions(name)

  def upload(self, name, version, target_dir, extra_tags=None):
    self._cipd.upload(name, version, target_dir, extra_tags=extra_tags)
    self._gs.upload(name, version, target_dir, extra_tags=extra_tags)

  def download(self, name, version, target_dir):
    self._gs.download(name, version, target_dir)

  def delete_contents(self, name):
    self._cipd.delete_contents(name)
    self._gs.delete_contents(name)


def _prompt(prompt):
  """Prompt for input, return result."""
  return raw_input(prompt)


class Asset(object):
  def __init__(self, name, store):
    self._store = store
    self._name = name
    self._dir = os.path.join(ASSETS_DIR, self._name)

  @property
  def version_file(self):
    """Return the path to the version file for this asset."""
    return os.path.join(self._dir, VERSION_FILENAME)

  def get_current_version(self):
    """Obtain the current version of the asset."""
    if not os.path.isfile(self.version_file):
      return -1
    with open(self.version_file) as f:
      return int(f.read())

  def get_available_versions(self):
    """Return the existing version numbers for this asset."""
    return self._store.get_available_versions(self._name)

  def get_next_version(self):
    """Find the next available version number for the asset."""
    versions = self.get_available_versions()
    if len(versions) == 0:
      return 0
    return versions[-1] + 1

  def download_version(self, version, target_dir):
    """Download the specified version of the asset."""
    self._store.download(self._name, version, target_dir)

  def download_current_version(self, target_dir):
    """Download the version of the asset specified in its version file."""
    v = self.get_current_version()
    self.download_version(v, target_dir)

  def upload_new_version(self, target_dir, commit=False, extra_tags=None):
    """Upload a new version and update the version file for the asset."""
    version = self.get_next_version()
    self._store.upload(self._name, version, target_dir, extra_tags=extra_tags)

    def _write_version():
      with open(self.version_file, 'w') as f:
        f.write(str(version))
      subprocess.check_call([utils.GIT, 'add', self.version_file])

    with utils.chdir(SKIA_DIR):
      if commit:
        with utils.git_branch():
          _write_version()
          subprocess.check_call([
              utils.GIT, 'commit', '-m', 'Update %s version' % self._name])
          subprocess.check_call([utils.GIT, 'cl', 'upload', '--bypass-hooks'])
      else:
        _write_version()

  @classmethod
  def add(cls, name, store):
    """Add an asset."""
    asset = cls(name, store)
    if os.path.isdir(asset._dir):
      raise Exception('Asset %s already exists!' % asset._name)

    print 'Creating asset in %s' % asset._dir
    os.mkdir(asset._dir)
    def copy_script(script):
      src = os.path.join(ASSETS_DIR, 'scripts', script)
      dst = os.path.join(asset._dir, script)
      print 'Creating %s' % dst
      shutil.copy(src, dst)
      subprocess.check_call([utils.GIT, 'add', dst])

    for script in ('download.py', 'upload.py', 'common.py'):
      copy_script(script)
    resp = _prompt('Add script to automate creation of this asset? (y/n) ')
    if resp == 'y':
      copy_script('create.py')
      copy_script('create_and_upload.py')
      print 'You will need to add implementation to the creation script.'
    print 'Successfully created asset %s.' % asset._name
    return asset

  def remove(self, remove_in_store=False):
    """Remove this asset."""
    # Ensure that the asset exists.
    if not os.path.isdir(self._dir):
      raise Exception('Asset %s does not exist!' % self._name)

    # Cleanup the store.
    if remove_in_store:
      self._store.delete_contents(self._name)

    # Remove the asset.
    subprocess.check_call([utils.GIT, 'rm', '-rf', self._dir])
    if os.path.isdir(self._dir):
      shutil.rmtree(self._dir)