jumpserver/apps/assets/migrations/0103_auto_20220902_1021.py

104 lines
3.5 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
node_model = apps.get_model('assets', 'Node')
org_model = apps.get_model('orgs', 'Organization')
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()
org_ids = gateways.order_by('org_id').values_list('org_id', flat=True).distinct()
node_dict = {}
for org_id in org_ids:
org = org_model.objects.using(db_alias).filter(id=org_id).first()
node = node_model.objects.using(db_alias).filter(
org_id=org_id, value=org.name, full_value=f'/{org.name}'
).first()
node_dict[org_id] = node
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)
node = node_dict.get(str(gateway.org_id))
asset.nodes.set([node])
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('accounts', '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', '0102_auto_20220816_1022'),
]
operations = [
migrations.RunPython(migrate_gateway_to_asset),
migrations.DeleteModel(
name='Gateway',
),
migrations.CreateModel(
name='Gateway',
fields=[
],
options={
'proxy': True,
'indexes': [],
'constraints': [],
'verbose_name': 'Gateway'
},
bases=('assets.host',),
),
]