quay/util/migrate/migrategithubdeploykeys.py

import logging
import logging.config
import json

from data.database import RepositoryBuildTrigger, BuildTriggerService, db, db_for_update
from app import app
from endpoints.trigger import BuildTriggerHandler
from util.security.ssh import generate_ssh_keypair
from github import GithubException

logger = logging.getLogger(__name__)

def backfill_github_deploykeys():
  """ Generates and saves private deploy keys for any GitHub build triggers still relying on
      the old buildpack behavior. """
  logger.setLevel(logging.DEBUG)
  logger.debug('GitHub deploy key backfill: Began execution')

  encountered = set()
  github_service = BuildTriggerService.get(name='github')

  while True:
    build_trigger_ids = list(RepositoryBuildTrigger
                              .select(RepositoryBuildTrigger.id)
                              .where(RepositoryBuildTrigger.private_key >> None)
                              .where(RepositoryBuildTrigger.service == github_service)
                              .where(RepositoryBuildTrigger.used_legacy_github >> None)
                              .limit(100))

    filtered_ids = [trigger.id for trigger in build_trigger_ids if trigger.id not in encountered]
    if len(filtered_ids) == 0:
      # We're done!
      logger.debug('GitHub deploy key backfill: Backfill completed')
      return

    logger.debug('GitHub deploy key backfill: Found %s records to update', len(filtered_ids))
    for trigger_id in filtered_ids:
      encountered.add(trigger_id)
      logger.debug('Updating build trigger: %s', trigger_id)

      with app.config['DB_TRANSACTION_FACTORY'](db):
        try:
          query = RepositoryBuildTrigger.select().where(RepositoryBuildTrigger.id == trigger_id)
          trigger = db_for_update(query).get()
        except RepositoryBuildTrigger.DoesNotExist:
          logger.debug('Could not find build trigger %s', trigger_id)
          continue

        trigger.used_legacy_github = True
        trigger.save()

        handler = BuildTriggerHandler.get_handler(trigger)

        config = handler.config
        if not 'build_source' in config:
          logger.debug('Could not find build source for trigger %s', trigger_id)
          continue

        build_source = config['build_source']
        gh_client = handler._get_client()

        # Find the GitHub repository.
        try:
          gh_repo = gh_client.get_repo(build_source)
        except GithubException:
          logger.exception('Cannot find repository %s for trigger %s', build_source, trigger.id)
          continue

        # Add a deploy key to the GitHub repository.
        public_key, private_key = generate_ssh_keypair()
        config['credentials'] = [
          {
            'name': 'SSH Public Key',
            'value': public_key,
          },
        ]

        logger.debug('Adding deploy key to build trigger %s', trigger.id)
        try:
          deploy_key = gh_repo.create_key('%s Builder' % app.config['REGISTRY_TITLE'], public_key)
          config['deploy_key_id'] = deploy_key.id
        except GithubException:
          logger.exception('Cannot add deploy key to repository %s for trigger %s', build_source, trigger.id)
          continue

        logger.debug('Saving deploy key for trigger %s', trigger.id)
        trigger.used_legacy_github = True
        trigger.private_key = private_key
        trigger.config = json.dumps(config)
        trigger.save()


if __name__ == "__main__":
  logging.getLogger('boto').setLevel(logging.CRITICAL)
  logging.getLogger('github').setLevel(logging.CRITICAL)

  logging.config.fileConfig('conf/logging_debug.conf', disable_existing_loggers=False)
  backfill_github_deploykeys()