jumpserver/apps/assets/migrations/0100_auto_20220711_1413.py

# Generated by Django 3.2.12 on 2022-07-11 06:13

import time
from django.utils import timezone
from itertools import groupby
from django.db import migrations


def migrate_asset_accounts(apps, schema_editor):
    auth_book_model = apps.get_model('assets', 'AuthBook')
    account_model = apps.get_model('accounts', 'Account')
    account_history_model = apps.get_model('accounts', 'HistoricalAccount')

    count = 0
    bulk_size = 1000
    print("\n\tStart migrate asset accounts")
    while True:
        start = time.time()
        auth_books = auth_book_model.objects \
                         .prefetch_related('systemuser') \
                         .all()[count:count + bulk_size]
        if not auth_books:
            break

        count += len(auth_books)
        # auth book 和 account 相同的属性
        same_attrs = [
            'id', 'username', 'comment', 'date_created', 'date_updated',
            'created_by', 'asset_id', 'org_id',
        ]
        # 认证的属性，可能是 auth_book 的，可能是 system_user 的
        auth_attrs = ['password', 'private_key', 'token']
        all_attrs = same_attrs + auth_attrs

        accounts = []
        for auth_book in auth_books:
            account_values = {'version': 1}

            system_user = auth_book.systemuser
            if system_user:
                # 更新一次系统用户的认证属性
                account_values.update({attr: getattr(system_user, attr, '') for attr in all_attrs})
                account_values['created_by'] = str(system_user.id)
                account_values['privileged'] = system_user.type == 'admin' \
                                                or system_user.username in ['root', 'Administrator']

            auth_book_auth = {attr: getattr(auth_book, attr, '') for attr in all_attrs if getattr(auth_book, attr, '')}
            # 最终优先使用 auth_book 的认证属性
            account_values.update(auth_book_auth)

            auth_infos = []
            username = account_values.get('username')
            if not username:
                continue

            for attr in auth_attrs:
                secret = account_values.pop(attr, None)
                if not secret:
                    continue

                if attr == 'private_key':
                    secret_type = 'ssh_key'
                    name = f'{username}(ssh key)'
                elif attr == 'token':
                    secret_type = 'token'
                    name = f'{username}(token)'
                else:
                    secret_type = attr
                    name = username
                auth_infos.append((name, secret_type, secret))

            if not auth_infos:
                auth_infos.append((username, 'password', ''))

            for name, secret_type, secret in auth_infos:
                if not name:
                    continue
                account = account_model(**account_values, name=name, secret=secret, secret_type=secret_type)
                accounts.append(account)

        accounts.sort(key=lambda x: (x.name, x.asset_id, x.date_updated))
        grouped_accounts = groupby(accounts, lambda x: (x.name, x.asset_id))

        accounts_to_add = []
        accounts_to_history = []
        for key, _accounts in grouped_accounts:
            _accounts = list(_accounts)
            if not _accounts:
                continue
            _account = _accounts[-1]
            accounts_to_add.append(_account)
            _account_history = []

            for ac in _accounts:
                if not ac.secret:
                    continue
                if ac.id != _account.id and ac.secret == _account.secret:
                    continue
                history_data = {
                    'id':   _account.id,
                    'secret': ac.secret,
                    'secret_type': ac.secret_type,
                    'history_date': ac.date_updated,
                    'history_type': '~',
                    'history_change_reason': 'from account {}'.format(_account.name),
                }
                _account_history.append(account_history_model(**history_data))
            _account.version = len(_account_history)
            accounts_to_history.extend(_account_history)

        account_model.objects.bulk_create(accounts_to_add, ignore_conflicts=True)
        account_history_model.objects.bulk_create(accounts_to_history, ignore_conflicts=True)
        print("\t  - Create asset accounts: {}-{} using: {:.2f}s".format(
            count - len(auth_books), count, time.time() - start
        ))
        print("\t      - accounts: {}".format(len(accounts_to_add)))
        print("\t      - histories: {}".format(len(accounts_to_history)))


def migrate_db_accounts(apps, schema_editor):
    app_perm_model = apps.get_model('perms', 'ApplicationPermission')
    account_model = apps.get_model('accounts', 'Account')
    perms = app_perm_model.objects.filter(category__in=['db', 'cloud'])

    same_attrs = [
        'username', 'comment', 'date_created', 'date_updated',
        'created_by', 'org_id',
    ]
    auth_attrs = ['password', 'private_key', 'token']
    all_attrs = same_attrs + auth_attrs

    print("\n\tStart migrate app accounts")

    index = 0
    total = perms.count()

    for perm in perms:
        index += 1
        start = time.time()

        apps = perm.applications.all()
        system_users = perm.system_users.all()
        accounts = []
        for s in system_users:
            values = {'version': 1}
            values.update({attr: getattr(s, attr, '') for attr in all_attrs})
            values['created_by'] = str(s.id)

            auth_infos = []
            username = values['username']
            for attr in auth_attrs:
                secret = values.pop(attr, None)
                if not secret:
                    continue

                if attr == 'private_key':
                    secret_type = 'ssh_key'
                    name = f'{username}(ssh key)'
                elif attr == 'token':
                    secret_type = 'token'
                    name = f'{username}(token)'
                else:
                    secret_type = attr
                    name = username
                auth_infos.append((name, secret_type, secret))

            if not auth_infos:
                auth_infos.append((username, 'password', ''))

            for name, secret_type, secret in auth_infos:
                values['name'] = name
                values['secret_type'] = secret_type
                values['secret'] = secret

                if not name:
                    continue

                for app in apps:
                    values['asset_id'] = str(app.id)
                    account = account_model(**values)
                    accounts.append(account)

        account_model.objects.bulk_create(accounts, ignore_conflicts=True)

        print("\t  - Progress ({}/{}), Create app accounts: {} using: {:.2f}s".format(
            index, total, len(accounts), time.time() - start
        ))


class Migration(migrations.Migration):
    dependencies = [
        ('accounts', '0001_initial'),
        ('assets', '0099_auto_20220711_1409'),
    ]

    operations = [
        migrations.RunPython(migrate_asset_accounts),
        migrations.RunPython(migrate_db_accounts),
    ]