CTFd/migrations/1_2_0_upgrade_2_0_0.py

from __future__ import print_function
import sys
import os

# This is important to allow access to the CTFd application factory
sys.path.append(os.getcwd())

import datetime
import hashlib
import netaddr
from flask_sqlalchemy import SQLAlchemy
from passlib.hash import bcrypt_sha256
from sqlalchemy.sql.expression import union_all
from CTFd import config, create_app
from sqlalchemy_utils import (
    drop_database,
)
from six.moves import input
import dataset

if __name__ == '__main__':
    print("/*\\ Migrating your database to 2.0.0 can potentially lose data./*\\")
    print("""/*\\ Please be sure to back up all data by:
        * creating a CTFd export
        * creating a dump of your actual database
        * and backing up the CTFd source code directory""")
    print("/*\\ CTFd maintainers are not responsible for any data loss! /*\\")
    if input('Run database migrations (Y/N)').lower().strip() == 'y':
        pass
    else:
        print('/*\\ Aborting database migrations... /*\\')
        print('/*\\ Exiting... /*\\')
        exit(1)

    db_url = config.Config.SQLALCHEMY_DATABASE_URI
    old_data = {}

    old_conn = dataset.connect(config.Config.SQLALCHEMY_DATABASE_URI)
    tables = old_conn.tables
    for table in tables:
        old_data[table] = old_conn[table].all()

    if 'alembic_version' in old_data:
        old_data.pop('alembic_version')

    print('Current Tables:')
    for table in old_data.keys():
        print('\t', table)

    old_conn.executable.close()

    print('DROPPING DATABASE')
    drop_database(db_url)

    app = create_app()

    new_conn = dataset.connect(config.Config.SQLALCHEMY_DATABASE_URI)

    print('MIGRATING Challenges')
    for challenge in old_data['challenges']:
        hidden = challenge.pop('hidden')
        challenge['state'] = 'hidden' if hidden else 'visible'
        new_conn['challenges'].insert(dict(challenge))
    del old_data['challenges']

    print('MIGRATING Teams')
    for team in old_data['teams']:
        admin = team.pop('admin')
        team['type'] = 'admin' if admin else 'user'
        team['hidden'] = bool(team.pop('banned'))
        team['banned'] = False
        team['verified'] = bool(team.pop('verified'))
        new_conn['users'].insert(dict(team))
    del old_data['teams']

    print('MIGRATING Pages')
    for page in old_data['pages']:
        page['content'] = page.pop('html')
        new_conn['pages'].insert(dict(page))
    del old_data['pages']

    print('MIGRATING Keys')
    for key in old_data['keys']:
        key['challenge_id'] = key.pop('chal')
        key['content'] = key.pop('flag')
        new_conn['flags'].insert(dict(key))
    del old_data['keys']

    print('MIGRATING Tags')
    for tag in old_data['tags']:
        tag['challenge_id'] = tag.pop('chal')
        tag['value'] = tag.pop('tag')
        new_conn['tags'].insert(dict(tag))
    del old_data['tags']

    print('MIGRATING Files')
    for f in old_data['files']:
        challenge_id = f.pop('chal')
        if challenge_id:
            f['challenge_id'] = challenge_id
            f['type'] = 'challenge'
        else:
            f['page_id'] = None
            f['type'] = 'page'
        new_conn['files'].insert(dict(f))
    del old_data['files']

    print('MIGRATING Hints')
    for hint in old_data['hints']:
        hint['type'] = 'standard'
        hint['challenge_id'] = hint.pop('chal')
        hint['content'] = hint.pop('hint')
        new_conn['hints'].insert(dict(hint))
    del old_data['hints']

    print('MIGRATING Unlocks')
    for unlock in old_data['unlocks']:
        unlock['user_id'] = unlock.pop('teamid')  # This is intentional as previous CTFds are effectively in user mode
        unlock['target'] = unlock.pop('item_id')
        unlock['type'] = unlock.pop('model')
        new_conn['unlocks'].insert(dict(unlock))
    del old_data['unlocks']

    print('MIGRATING Awards')
    for award in old_data['awards']:
        award['user_id'] = award.pop('teamid')  # This is intentional as previous CTFds are effectively in user mode
        new_conn['awards'].insert(dict(award))
    del old_data['awards']

    submissions = []
    for solve in old_data['solves']:
        solve.pop('id')  # ID of a solve doesn't really matter
        solve['challenge_id'] = solve.pop('chalid')
        solve['user_id'] = solve.pop('teamid')
        solve['provided'] = solve.pop('flag')
        solve['type'] = 'correct'
        solve['model'] = 'solve'
        submissions.append(solve)

    for wrong_key in old_data['wrong_keys']:
        wrong_key.pop('id')  # ID of a fail doesn't really matter.
        wrong_key['challenge_id'] = wrong_key.pop('chalid')
        wrong_key['user_id'] = wrong_key.pop('teamid')
        wrong_key['provided'] = wrong_key.pop('flag')
        wrong_key['type'] = 'incorrect'
        wrong_key['model'] = 'wrong_key'
        submissions.append(wrong_key)

    submissions = sorted(submissions, key=lambda k: k['date'])
    print('MIGRATING Solves & WrongKeys')
    for submission in submissions:
        model = submission.pop('model')
        if model == 'solve':
            new_id = new_conn['submissions'].insert(dict(submission))
            submission['id'] = new_id
            new_conn['solves'].insert(dict(submission))
        elif model == 'wrong_key':
            new_conn['submissions'].insert(dict(submission))
    del old_data['solves']
    del old_data['wrong_keys']

    print('MIGRATING Tracking')
    for tracking in old_data['tracking']:
        tracking['user_id'] = tracking.pop('team')
        new_conn['tracking'].insert(dict(tracking))
    del old_data['tracking']

    print('MIGRATING Config')
    banned = [
        'ctf_version',
    ]
    workshop_mode = None
    hide_scores = None
    prevent_registration = None
    view_challenges_unregistered = None
    view_scoreboard_if_authed = None

    challenge_visibility = 'private'
    registration_visibility = 'public'
    score_visibility = 'public'
    account_visibility = 'public'
    for config in old_data['config']:
        config.pop('id')

        if config['key'] == 'workshop_mode':
            workshop_mode = config['value']
        elif config['key'] == 'hide_scores':
            hide_scores = config['value']
        elif config['key'] == 'prevent_registration':
            prevent_registration = config['value']
        elif config['key'] == 'view_challenges_unregistered':
            view_challenges_unregistered = config['value']
        elif config['key'] == 'view_scoreboard_if_authed':
            view_scoreboard_if_authed = config['value']

        if config['key'] not in banned:
            new_conn['config'].insert(dict(config))

    if workshop_mode:
        score_visibility = 'admins'
        account_visibility = 'admins'

    if hide_scores:
        score_visibility = 'hidden'

    if prevent_registration:
        registration_visibility = 'private'

    if view_challenges_unregistered:
        challenge_visibility = 'public'

    if view_scoreboard_if_authed:
        score_visibility = 'private'

    new_conn['config'].insert({
        'key': 'user_mode',
        'value': 'users'
    })
    new_conn['config'].insert({
        'key': 'challenge_visibility',
        'value': challenge_visibility
    })
    new_conn['config'].insert({
        'key': 'registration_visibility',
        'value': registration_visibility
    })
    new_conn['config'].insert({
        'key': 'score_visibility',
        'value': score_visibility
    })
    new_conn['config'].insert({
        'key': 'account_visibility',
        'value': account_visibility
    })
    del old_data['config']

    manual = []
    not_created = []
    print('MIGRATING extra tables')
    for table in old_data.keys():
        print('MIGRATING', table)
        new_conn.create_table(table, primary_id=False)
        data = old_data[table]

        ran = False
        for row in data:
            new_conn[table].insert(dict(row))
            ran = True
        else:  # We finished inserting
            if ran:
                manual.append(table)

        if ran is False:
            not_created.append(table)

    print('Migration completed.')
    print('The following tables require manual setting of primary keys and manual inspection')
    for table in manual:
        print('\t', table)

    print('For example you can use the following commands if you know that the PRIMARY KEY for the table is `id`:')
    for table in manual:
        print('\t', 'ALTER TABLE `{table}` ADD PRIMARY KEY(id)'.format(table=table))

    print(
        'The following tables were not created because they were empty and must be manually recreated (e.g. app.db.create_all()')
    for table in not_created:
        print('\t', table)