# Generated by Django 3.2.12 on 2022-07-11 06:13 import time import math 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 = [ '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['privileged'] = system_user.type == 'admin' \ or system_user.username in ['root', 'Administrator'] if system_user.su_enabled and system_user.su_from: created_by = f'{str(system_user.id)}::{str(system_user.su_from.username)}' else: created_by = str(system_user.id) account_values['created_by'] = created_by 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 update_asset_accounts_su_from(apps, schema_editor): # Update accounts su_from print("\n\tStart update asset accounts su_from field") account_model = apps.get_model('accounts', 'Account') platform_model = apps.get_model('assets', 'Platform') asset_model = apps.get_model('assets', 'Asset') platform_ids = list(platform_model.objects.filter(su_enabled=True).values_list('id', flat=True)) count = 0 step_size = 1000 count_account = 0 while True: start = time.time() asset_ids = asset_model.objects \ .filter(platform_id__in=platform_ids) \ .values_list('id', flat=True)[count:count + step_size] asset_ids = list(asset_ids) if not asset_ids: break count += len(asset_ids) accounts = list(account_model.objects.filter(asset_id__in=asset_ids)) # {asset_id_account_username: account.id}} asset_accounts_mapper = {} for a in accounts: try: k = f'{a.asset_id}_{a.username}' asset_accounts_mapper[k] = str(a.id) except Exception as e: pass update_accounts = [] for a in accounts: try: if not a.created_by: continue created_by_list = a.created_by.split('::') if len(created_by_list) != 2: continue su_from_username = created_by_list[1] if not su_from_username: continue k = f'{a.asset_id}_{su_from_username}' su_from_id = asset_accounts_mapper.get(k) if not su_from_id: continue a.su_from_id = su_from_id update_accounts.append(a) except Exception as e: pass count_account += len(update_accounts) log_msg = "\t - [{}]: Update accounts su_from: {}-{} {:.2f}s" try: account_model.objects.bulk_update(update_accounts, ['su_from_id']) except Exception as e: status = 'Failed' else: status = 'Success' print(log_msg.format(status, count_account - len(update_accounts), count_account, time.time() - start)) 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 or f'{username}(password)' auth_infos.append((name, secret_type, secret)) if not auth_infos: name = username or f'{username}(password)' auth_infos.append((name, '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(update_asset_accounts_su_from), migrations.RunPython(migrate_db_accounts), ]