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

197 lines
7.1 KiB
Python

# 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['username']
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),
]