jumpserver/apps/assets/migrations/0098_auto_20220430_2126.py

# Generated by Django 3.1.14 on 2022-04-26 07:58
import uuid

from django.db import migrations

failed_apps = []


def get_prop_name_id(apps, app, category):
    asset_model = apps.get_model('assets', 'Asset')
    _id = app.id
    id_exists = asset_model.objects.filter(id=_id).exists()
    if id_exists:
        _id = uuid.uuid4()
    name = app.name
    name_exists = asset_model.objects.filter(name=name).exists()
    if name_exists:
        name = category + '-' + app.name
    return _id, name


def migrate_database_to_asset(apps, *args):
    node_model = apps.get_model('assets', 'Node')
    app_model = apps.get_model('applications', 'Application')
    db_model = apps.get_model('assets', 'Database')
    platform_model = apps.get_model('assets', 'Platform')

    applications = app_model.objects.filter(category='db')
    platforms = platform_model.objects.all().filter(internal=True).exclude(name='Redis6+')
    platforms_map = {p.type: p for p in platforms}
    print()

    for app in applications:
        attrs = {'host': '', 'port': 0, 'database': ''}
        _attrs = app.attrs or {}
        attrs.update(_attrs)

        name = 'DB-{}'.format(app.name)
        db = db_model(
            id=app.id, name=name, address=attrs['host'],
            protocols='{}/{}'.format(app.type, attrs['port']),
            db_name=attrs['database'] or '',
            platform=platforms_map[app.type],
            org_id=app.org_id
        )
        try:
            print("\t- Create database: ", app.name)
            db.save()
        except:
            failed_apps.append(app)
            pass


def migrate_cloud_to_asset(apps, *args):
    app_model = apps.get_model('applications', 'Application')
    cloud_model = apps.get_model('assets', 'Cloud')
    platform_model = apps.get_model('assets', 'Platform')

    applications = app_model.objects.filter(category='cloud')
    platform = platform_model.objects.filter(type='k8s').first()
    print()

    for app in applications:
        attrs = app.attrs
        print("\t- Create cloud: {}".format(app.name))
        name = 'Cloud-{}'.format(app.name)
        cloud = cloud_model(
            id=app.id, name=name,
            address=attrs.get('cluster', ''),
            protocols='k8s/443', platform=platform,
            org_id=app.org_id,
        )

        try:
            cloud.save()
        except Exception as e:
            failed_apps.append(cloud)
            print("Error: ", e)


def create_app_nodes(apps, org_id):
    node_model = apps.get_model('assets', 'Node')

    child_pattern = r'^[0-9]+:[0-9]+$'
    node_keys = node_model.objects.filter(org_id=org_id) \
        .filter(key__regex=child_pattern) \
        .values_list('key', flat=True)
    if node_keys:
        node_key_split = [key.split(':') for key in node_keys]
        next_value = max([int(k[1]) for k in node_key_split]) + 1
        parent_key = node_key_split[0][0]
    else:
        root_node = node_model.objects.filter(org_id=org_id) \
            .filter(parent_key='', key__regex=r'^[0-9]+$') \
            .exclude(key__startswith='-') \
            .first()
        if not root_node:
            return
        parent_key = root_node.key
        next_value = 0

    next_key = '{}:{}'.format(parent_key, next_value)
    name = 'Apps'
    parent = node_model.objects.get(key=parent_key)
    full_value = parent.full_value + '/' + name
    defaults = {
        'key': next_key, 'value': name, 'parent_key': parent_key,
        'full_value': full_value, 'org_id': org_id
    }
    node, created = node_model.objects.get_or_create(
        defaults=defaults, value=name, org_id=org_id,
    )
    node.parent = parent
    return node


def migrate_to_nodes(apps, *args):
    org_model = apps.get_model('orgs', 'Organization')
    asset_model = apps.get_model('assets', 'Asset')
    orgs = org_model.objects.all()

    # Todo: 优化一些
    for org in orgs:
        node = create_app_nodes(apps, org.id)
        assets = asset_model.objects.filter(
            platform__category__in=['remote_app', 'database', 'cloud'],
            org_id=org.id
        )
        if not node:
            continue
        print("\t- Set node asset: ", node)
        node.assets_amount = len(assets)
        node.save()
        node.assets.set(assets)
        parent = node.parent
        parent.assets_amount += len(assets)
        parent.save()


def migrate_ori_host_to_devices(apps, *args):
    device_model = apps.get_model('assets', 'Device')
    asset_model = apps.get_model('assets', 'Asset')
    host_model = apps.get_model('assets', 'Host')
    hosts_need_migrate_to_device = host_model.objects.filter(asset_ptr__platform__category='device')
    assets = asset_model.objects.filter(id__in=hosts_need_migrate_to_device.values_list('asset_ptr_id', flat=True))
    assets_map = {asset.id: asset for asset in assets}

    print("\t- Migrate ori host to device: ", len(hosts_need_migrate_to_device))
    for host in hosts_need_migrate_to_device:
        asset = assets_map.get(host.asset_ptr_id)
        if not asset:
            continue
        device = device_model(asset_ptr_id=asset.id)
        device.__dict__.update(asset.__dict__)
        device.save()
        host.delete(keep_parents=True)


class Migration(migrations.Migration):
    dependencies = [
        ('assets', '0097_auto_20220426_1558'),
        ('applications', '0020_auto_20220316_2028')
    ]

    operations = [
        migrations.RunPython(migrate_database_to_asset),
        migrations.RunPython(migrate_cloud_to_asset),
        migrations.RunPython(migrate_to_nodes),
        migrations.RunPython(migrate_ori_host_to_devices),
    ]