jumpserver/apps/assets/migrations/0112_gateway_to_asset.py

74 lines
2.6 KiB
Python

# Generated by Django 3.2.13 on 2022-09-29 11:03
from django.db import migrations
from assets.const.host import GATEWAY_NAME
def _create_account_obj(secret, secret_type, gateway, asset, account_model):
return account_model(
asset=asset,
secret=secret,
org_id=gateway.org_id,
secret_type=secret_type,
username=gateway.username,
name=f'{gateway.name}-{secret_type}-{GATEWAY_NAME.lower()}',
)
def migrate_gateway_to_asset(apps, schema_editor):
db_alias = schema_editor.connection.alias
gateway_model = apps.get_model('assets', 'Gateway')
platform_model = apps.get_model('assets', 'Platform')
gateway_platform = platform_model.objects.using(db_alias).get(name=GATEWAY_NAME)
print('>>> migrate gateway to asset')
asset_dict = {}
host_model = apps.get_model('assets', 'Host')
asset_model = apps.get_model('assets', 'Asset')
protocol_model = apps.get_model('assets', 'Protocol')
gateways = gateway_model.objects.all()
for gateway in gateways:
comment = gateway.comment if gateway.comment else ''
data = {
'comment': comment,
'name': f'{gateway.name}-{GATEWAY_NAME.lower()}',
'address': gateway.ip,
'domain': gateway.domain,
'org_id': gateway.org_id,
'is_active': gateway.is_active,
'platform': gateway_platform,
}
asset = asset_model.objects.using(db_alias).create(**data)
asset_dict[gateway.id] = asset
protocol_model.objects.using(db_alias).create(name='ssh', port=gateway.port, asset=asset)
hosts = [host_model(asset_ptr=asset) for asset in asset_dict.values()]
host_model.objects.using(db_alias).bulk_create(hosts, ignore_conflicts=True)
print('>>> migrate gateway to account')
accounts = []
account_model = apps.get_model('assets', 'Account')
for gateway in gateways:
password = gateway.password
private_key = gateway.private_key
asset = asset_dict[gateway.id]
if password:
accounts.append(_create_account_obj(
password, 'password', gateway, asset, account_model
))
if private_key:
accounts.append(_create_account_obj(
private_key, 'ssh_key', gateway, asset, account_model
))
account_model.objects.using(db_alias).bulk_create(accounts)
class Migration(migrations.Migration):
dependencies = [
('assets', '0111_alter_automationexecution_status'),
]
operations = [
migrations.RunPython(migrate_gateway_to_asset),
]