This repository has been archived on 2020-03-24. You can view files and clone it, but cannot push or open issues or pull requests.
quay/data/migrations/versions/703298a825c2_backfill_new_encrypted_fields.py
Joseph Schorr a54fb1b23a Fix the encrypted token migration issue encountered on HEAD
This change ensures there is better messaging around the encrypted token migration, including a new phase to use for new installations, and fixes an issue encountered when running database migrations for new installations
2019-11-14 14:50:47 -05:00

287 lines
12 KiB
Python

"""Backfill new encrypted fields
Revision ID: 703298a825c2
Revises: c13c8052f7a6
Create Date: 2019-08-19 16:07:48.109889
"""
# revision identifiers, used by Alembic.
revision = '703298a825c2'
down_revision = 'c13c8052f7a6'
import logging
import uuid
from datetime import datetime
from peewee import (JOIN, IntegrityError, DateTimeField, CharField, ForeignKeyField,
BooleanField, TextField, IntegerField)
from alembic import op as original_op
from data.migrations.progress import ProgressWrapper
import sqlalchemy as sa
from data.database import (BaseModel, User, Repository, AccessTokenKind, Role,
random_string_generator, QuayUserField, BuildTriggerService,
uuid_generator, DisableReason)
from data.fields import Credential, DecryptedValue, EncryptedCharField, EncryptedTextField, EnumField, CredentialField
from data.model.token import ACCESS_TOKEN_NAME_PREFIX_LENGTH
from data.model.appspecifictoken import TOKEN_NAME_PREFIX_LENGTH as AST_TOKEN_NAME_PREFIX_LENGTH
from data.model.oauth import ACCESS_TOKEN_PREFIX_LENGTH as OAUTH_ACCESS_TOKEN_PREFIX_LENGTH
from data.model.oauth import AUTHORIZATION_CODE_PREFIX_LENGTH
BATCH_SIZE = 10
logger = logging.getLogger(__name__)
def _iterate(model_class, clause):
while True:
has_rows = False
for row in list(model_class.select().where(clause).limit(BATCH_SIZE)):
has_rows = True
yield row
if not has_rows:
break
def _decrypted(value):
if value is None:
return None
assert isinstance(value, basestring)
return DecryptedValue(value)
# NOTE: As per standard migrations involving Peewee models, we copy them here, as they will change
# after this call.
class AccessToken(BaseModel):
code = CharField(default=random_string_generator(length=64), unique=True, index=True)
token_name = CharField(default=random_string_generator(length=32), unique=True, index=True)
token_code = EncryptedCharField(default_token_length=32)
class RobotAccountToken(BaseModel):
robot_account = QuayUserField(index=True, allows_robots=True, unique=True)
token = EncryptedCharField(default_token_length=64)
fully_migrated = BooleanField(default=False)
class RepositoryBuildTrigger(BaseModel):
uuid = CharField(default=uuid_generator, index=True)
auth_token = CharField(null=True)
private_key = TextField(null=True)
secure_auth_token = EncryptedCharField(null=True)
secure_private_key = EncryptedTextField(null=True)
fully_migrated = BooleanField(default=False)
class AppSpecificAuthToken(BaseModel):
token_name = CharField(index=True, unique=True, default=random_string_generator(60))
token_secret = EncryptedCharField(default_token_length=60)
token_code = CharField(default=random_string_generator(length=120), unique=True, index=True)
class OAuthAccessToken(BaseModel):
token_name = CharField(index=True, unique=True)
token_code = CredentialField()
access_token = CharField(index=True)
class OAuthAuthorizationCode(BaseModel):
code = CharField(index=True, unique=True, null=True)
code_name = CharField(index=True, unique=True)
code_credential = CredentialField()
class OAuthApplication(BaseModel):
secure_client_secret = EncryptedCharField(default_token_length=40, null=True)
fully_migrated = BooleanField(default=False)
client_secret = CharField(default=random_string_generator(length=40))
def upgrade(tables, tester, progress_reporter):
op = ProgressWrapper(original_op, progress_reporter)
from app import app
if app.config.get('SETUP_COMPLETE', False) or tester.is_testing:
# Empty all access token names to fix the bug where we put the wrong name and code
# in for some tokens.
AccessToken.update(token_name=None).where(AccessToken.token_name >> None).execute()
# AccessToken.
logger.info('Backfilling encrypted credentials for access tokens')
for access_token in _iterate(AccessToken, ((AccessToken.token_name >> None) |
(AccessToken.token_name == ''))):
logger.info('Backfilling encrypted credentials for access token %s', access_token.id)
assert access_token.code is not None
assert access_token.code[:ACCESS_TOKEN_NAME_PREFIX_LENGTH]
assert access_token.code[ACCESS_TOKEN_NAME_PREFIX_LENGTH:]
token_name = access_token.code[:ACCESS_TOKEN_NAME_PREFIX_LENGTH]
token_code = _decrypted(access_token.code[ACCESS_TOKEN_NAME_PREFIX_LENGTH:])
(AccessToken
.update(token_name=token_name, token_code=token_code)
.where(AccessToken.id == access_token.id, AccessToken.code == access_token.code)
.execute())
assert AccessToken.select().where(AccessToken.token_name >> None).count() == 0
# Robots.
logger.info('Backfilling encrypted credentials for robots')
while True:
has_row = False
query = (User
.select()
.join(RobotAccountToken, JOIN.LEFT_OUTER)
.where(User.robot == True, RobotAccountToken.id >> None)
.limit(BATCH_SIZE))
for robot_user in query:
logger.info('Backfilling encrypted credentials for robot %s', robot_user.id)
has_row = True
try:
RobotAccountToken.create(robot_account=robot_user,
token=_decrypted(robot_user.email),
fully_migrated=False)
except IntegrityError:
break
if not has_row:
break
# RepositoryBuildTrigger
logger.info('Backfilling encrypted credentials for repo build triggers')
for repo_build_trigger in _iterate(RepositoryBuildTrigger,
(RepositoryBuildTrigger.fully_migrated == False)):
logger.info('Backfilling encrypted credentials for repo build trigger %s',
repo_build_trigger.id)
(RepositoryBuildTrigger
.update(secure_auth_token=_decrypted(repo_build_trigger.auth_token),
secure_private_key=_decrypted(repo_build_trigger.private_key),
fully_migrated=True)
.where(RepositoryBuildTrigger.id == repo_build_trigger.id,
RepositoryBuildTrigger.uuid == repo_build_trigger.uuid)
.execute())
assert (RepositoryBuildTrigger
.select()
.where(RepositoryBuildTrigger.fully_migrated == False)
.count()) == 0
# AppSpecificAuthToken
logger.info('Backfilling encrypted credentials for app specific auth tokens')
for token in _iterate(AppSpecificAuthToken, ((AppSpecificAuthToken.token_name >> None) |
(AppSpecificAuthToken.token_name == '') |
(AppSpecificAuthToken.token_secret >> None))):
logger.info('Backfilling encrypted credentials for app specific auth %s',
token.id)
assert token.token_code[AST_TOKEN_NAME_PREFIX_LENGTH:]
token_name = token.token_code[:AST_TOKEN_NAME_PREFIX_LENGTH]
token_secret = _decrypted(token.token_code[AST_TOKEN_NAME_PREFIX_LENGTH:])
assert token_name
assert token_secret
(AppSpecificAuthToken
.update(token_name=token_name,
token_secret=token_secret)
.where(AppSpecificAuthToken.id == token.id,
AppSpecificAuthToken.token_code == token.token_code)
.execute())
assert (AppSpecificAuthToken
.select()
.where(AppSpecificAuthToken.token_name >> None)
.count()) == 0
# OAuthAccessToken
logger.info('Backfilling credentials for OAuth access tokens')
for token in _iterate(OAuthAccessToken, ((OAuthAccessToken.token_name >> None) |
(OAuthAccessToken.token_name == ''))):
logger.info('Backfilling credentials for OAuth access token %s', token.id)
token_name = token.access_token[:OAUTH_ACCESS_TOKEN_PREFIX_LENGTH]
token_code = Credential.from_string(token.access_token[OAUTH_ACCESS_TOKEN_PREFIX_LENGTH:])
assert token_name
assert token.access_token[OAUTH_ACCESS_TOKEN_PREFIX_LENGTH:]
(OAuthAccessToken
.update(token_name=token_name,
token_code=token_code)
.where(OAuthAccessToken.id == token.id,
OAuthAccessToken.access_token == token.access_token)
.execute())
assert (OAuthAccessToken
.select()
.where(OAuthAccessToken.token_name >> None)
.count()) == 0
# OAuthAuthorizationCode
logger.info('Backfilling credentials for OAuth auth code')
for code in _iterate(OAuthAuthorizationCode, ((OAuthAuthorizationCode.code_name >> None) |
(OAuthAuthorizationCode.code_name == ''))):
logger.info('Backfilling credentials for OAuth auth code %s', code.id)
user_code = code.code or random_string_generator(AUTHORIZATION_CODE_PREFIX_LENGTH * 2)()
code_name = user_code[:AUTHORIZATION_CODE_PREFIX_LENGTH]
code_credential = Credential.from_string(user_code[AUTHORIZATION_CODE_PREFIX_LENGTH:])
assert code_name
assert user_code[AUTHORIZATION_CODE_PREFIX_LENGTH:]
(OAuthAuthorizationCode
.update(code_name=code_name, code_credential=code_credential)
.where(OAuthAuthorizationCode.id == code.id)
.execute())
assert (OAuthAuthorizationCode
.select()
.where(OAuthAuthorizationCode.code_name >> None)
.count()) == 0
# OAuthApplication
logger.info('Backfilling secret for OAuth applications')
for app in _iterate(OAuthApplication, OAuthApplication.fully_migrated == False):
logger.info('Backfilling secret for OAuth application %s', app.id)
client_secret = app.client_secret or str(uuid.uuid4())
secure_client_secret = _decrypted(client_secret)
(OAuthApplication
.update(secure_client_secret=secure_client_secret, fully_migrated=True)
.where(OAuthApplication.id == app.id, OAuthApplication.fully_migrated == False)
.execute())
assert (OAuthApplication
.select()
.where(OAuthApplication.fully_migrated == False)
.count()) == 0
# Adjust existing fields to be nullable.
op.alter_column('accesstoken', 'code', nullable=True, existing_type=sa.String(length=255))
op.alter_column('oauthaccesstoken', 'access_token', nullable=True, existing_type=sa.String(length=255))
op.alter_column('oauthauthorizationcode', 'code', nullable=True, existing_type=sa.String(length=255))
op.alter_column('appspecificauthtoken', 'token_code', nullable=True, existing_type=sa.String(length=255))
# Adjust new fields to be non-nullable.
op.alter_column('accesstoken', 'token_name', nullable=False, existing_type=sa.String(length=255))
op.alter_column('accesstoken', 'token_code', nullable=False, existing_type=sa.String(length=255))
op.alter_column('appspecificauthtoken', 'token_name', nullable=False, existing_type=sa.String(length=255))
op.alter_column('appspecificauthtoken', 'token_secret', nullable=False, existing_type=sa.String(length=255))
op.alter_column('oauthaccesstoken', 'token_name', nullable=False, existing_type=sa.String(length=255))
op.alter_column('oauthaccesstoken', 'token_code', nullable=False, existing_type=sa.String(length=255))
op.alter_column('oauthauthorizationcode', 'code_name', nullable=False, existing_type=sa.String(length=255))
op.alter_column('oauthauthorizationcode', 'code_credential', nullable=False, existing_type=sa.String(length=255))
def downgrade(tables, tester, progress_reporter):
op = ProgressWrapper(original_op, progress_reporter)
op.alter_column('accesstoken', 'token_name', nullable=True, existing_type=sa.String(length=255))
op.alter_column('accesstoken', 'token_code', nullable=True, existing_type=sa.String(length=255))
op.alter_column('appspecificauthtoken', 'token_name', nullable=True, existing_type=sa.String(length=255))
op.alter_column('appspecificauthtoken', 'token_secret', nullable=True, existing_type=sa.String(length=255))
op.alter_column('oauthaccesstoken', 'token_name', nullable=True, existing_type=sa.String(length=255))
op.alter_column('oauthaccesstoken', 'token_code', nullable=True, existing_type=sa.String(length=255))
op.alter_column('oauthauthorizationcode', 'code_name', nullable=True, existing_type=sa.String(length=255))
op.alter_column('oauthauthorizationcode', 'code_credential', nullable=True, existing_type=sa.String(length=255))