mirror of https://github.com/jumpserver/jumpserver
271 lines
9.8 KiB
Python
271 lines
9.8 KiB
Python
# 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),
|
|
]
|