perf: 优化生成假数据 (#4759)

* perf: 优化生成假数据

utils/generate_fake_data/generate.py

@@ -0,0 +1,49 @@
+#!/usr/bin/env python
+#
+import os
+import sys
+import django
+import argparse
+
+BASE_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+APPS_DIR = os.path.join(BASE_DIR, 'apps')
+sys.path.insert(0, APPS_DIR)
+
+os.environ.setdefault("DJANGO_SETTINGS_MODULE", "jumpserver.settings")
+django.setup()
+
+from resources.assets import AssetsGenerator, NodesGenerator, SystemUsersGenerator, AdminUsersGenerator
+from resources.users import UserGroupGenerator, UserGenerator
+from resources.perms import AssetPermissionGenerator
+
+
+resource_generator_mapper = {
+    'asset': AssetsGenerator,
+    'node': NodesGenerator,
+    'system_user': SystemUsersGenerator,
+    'admin_user': AdminUsersGenerator,
+    'user': UserGenerator,
+    'user_group': UserGroupGenerator,
+    'asset_permission': AssetPermissionGenerator
+}
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Generate fake data')
+    parser.add_argument(
+        'resource', type=str,
+        choices=resource_generator_mapper.keys(),
+        help="resource to generate"
+    )
+    parser.add_argument('-c', '--count', type=int, default=100)
+    parser.add_argument('-b', '--batch_size', type=int, default=100)
+    parser.add_argument('-o', '--org', type=str, default='')
+    args = parser.parse_args()
+    resource, count, batch_size, org_id = args.resource, args.count, args.batch_size, args.org
+    generator_cls = resource_generator_mapper[resource]
+    generator = generator_cls(org_id=org_id, batch_size=batch_size)
+    generator.generate(count)
+
+
+if __name__ == '__main__':
+    main()

utils/generate_fake_data/resources/assets.py

@@ -0,0 +1,96 @@
+from random import choice
+import random
+import forgery_py
+
+from .base import FakeDataGenerator
+
+from assets.models import *
+from assets.utils import check_node_assets_amount
+
+
+class AdminUsersGenerator(FakeDataGenerator):
+    resource = 'admin_user'
+
+    def do_generate(self, batch, batch_size):
+        admin_users = []
+        for i in batch:
+            username = forgery_py.internet.user_name(True)
+            password = forgery_py.basic.password()
+            admin_users.append(AdminUser(
+                name=username.title(),
+                username=username,
+                password=password,
+                org_id=self.org.id,
+                created_by='Fake',
+            ))
+        AdminUser.objects.bulk_create(admin_users, ignore_conflicts=True)
+
+
+class SystemUsersGenerator(FakeDataGenerator):
+    def do_generate(self, batch, batch_size):
+        system_users = []
+        protocols = list(dict(SystemUser.PROTOCOL_CHOICES).keys())
+        for i in batch:
+            username = forgery_py.internet.user_name(True)
+            protocol = random.choice(protocols)
+            name = username.title()
+            name = f'{name}-{protocol}'
+            system_users.append(SystemUser(
+                name=name,
+                username=username,
+                password=forgery_py.basic.password(),
+                protocol=protocol,
+                org_id=self.org.id,
+                created_by='Fake',
+            ))
+        SystemUser.objects.bulk_create(system_users, ignore_conflicts=True)
+
+
+class NodesGenerator(FakeDataGenerator):
+    resource = 'node'
+
+    def do_generate(self, batch, batch_size):
+        nodes_to_generate_children = list(Node.objects.all())
+        for i in batch:
+            parent = random.choice(nodes_to_generate_children)
+            parent.create_child()
+
+
+class AssetsGenerator(FakeDataGenerator):
+    resource = 'asset'
+    admin_users_id: list
+    nodes_id: list
+
+    def pre_generate(self):
+        self.admin_users_id = list(AdminUser.objects.all().values_list('id', flat=True))
+        self.nodes_id = list(Node.objects.all().values_list('id', flat=True))
+
+    def set_assets_nodes(self, assets):
+        assets_id = [asset.id for asset in assets]
+        objs = []
+        for asset_id in assets_id:
+            nodes_id_add_to = random.sample(self.nodes_id, 3)
+            objs_add = [Asset.nodes.through(asset_id=asset_id, node_id=nid) for nid in nodes_id_add_to]
+            objs.extend(objs_add)
+        Asset.nodes.through.objects.bulk_create(objs, ignore_conflicts=True)
+
+    def do_generate(self, batch, batch_size):
+        assets = []
+
+        for i in batch:
+            ip = forgery_py.internet.ip_v4()
+            hostname = forgery_py.email.address().replace('@', '.')
+            hostname = f'{hostname}-{ip}'
+            data = dict(
+                ip=ip,
+                hostname=hostname,
+                admin_user_id=choice(self.admin_users_id),
+                created_by='Fake',
+                org_id=self.org.id
+            )
+            assets.append(Asset(**data))
+        creates = Asset.objects.bulk_create(assets, ignore_conflicts=True)
+        self.set_assets_nodes(creates)
+
+    def after_generate(self):
+        check_node_assets_amount()

utils/generate_fake_data/resources/base.py

@@ -0,0 +1,45 @@
+#!/usr/bin/python
+
+from random import seed
+from itertools import islice
+
+from orgs.models import Organization
+
+
+class FakeDataGenerator:
+    resource = 'Fake'
+
+    def __init__(self, batch_size=100, org_id=None):
+        self.batch_size = batch_size
+        self.org = self.switch_org(org_id)
+        seed()
+
+    def switch_org(self, org_id):
+        o = Organization.get_instance(org_id, default=True)
+        if o:
+            o.change_to()
+        print('Current org is: {}'.format(o))
+        return o
+
+    def do_generate(self, batch, batch_size):
+        raise NotImplementedError
+
+    def pre_generate(self):
+        pass
+
+    def after_generate(self):
+        pass
+
+    def generate(self, count=100):
+        self.pre_generate()
+        counter = iter(range(count))
+        created = 0
+        while True:
+            batch = list(islice(counter, self.batch_size))
+            if not batch:
+                break
+            self.do_generate(batch, self.batch_size)
+            from_size = created
+            created += len(batch)
+            print('Generate %s: %s-%s' % (self.resource, from_size, created))
+        self.after_generate()

utils/generate_fake_data/resources/perms.py

@@ -0,0 +1,82 @@
+from random import choice, sample
+import forgery_py
+
+from .base import FakeDataGenerator
+
+from users.models import *
+from assets.models import *
+from perms.models import *
+
+
+class AssetPermissionGenerator(FakeDataGenerator):
+    resource = 'asset_permission'
+    users_id: list
+    user_groups_id: list
+    assets_id: list
+    nodes_id: list
+    system_users_id: list
+
+    def pre_generate(self):
+        self.nodes_id = list(Node.objects.all().values_list('id', flat=True))
+        self.assets_id = list(Asset.objects.all().values_list('id', flat=True))
+        self.system_users_id = list(SystemUser.objects.all().values_list('id', flat=True))
+        self.users_id = list(User.objects.all().values_list('id', flat=True))
+        self.user_groups_id = list(UserGroup.objects.all().values_list('id', flat=True))
+
+    def set_users(self, perms):
+        through = AssetPermission.users.through
+        choices = self.users_id
+        relation_name = 'user_id'
+        self.set_relations(perms, through, relation_name, choices)
+
+    def set_user_groups(self, perms):
+        through = AssetPermission.user_groups.through
+        choices = self.user_groups_id
+        relation_name = 'usergroup_id'
+        self.set_relations(perms, through, relation_name, choices)
+
+    def set_assets(self, perms):
+        through = AssetPermission.assets.through
+        choices = self.assets_id
+        relation_name = 'asset_id'
+        self.set_relations(perms, through, relation_name, choices)
+
+    def set_nodes(self, perms):
+        through = AssetPermission.nodes.through
+        choices = self.nodes_id
+        relation_name = 'node_id'
+        self.set_relations(perms, through, relation_name, choices)
+
+    def set_system_users(self, perms):
+        through = AssetPermission.system_users.through
+        choices = self.system_users_id
+        relation_name = 'systemuser_id'
+        self.set_relations(perms, through, relation_name, choices)
+
+    def set_relations(self, perms, through, relation_name, choices, choice_count=None):
+        relations = []
+
+        for perm in perms:
+            if choice_count is None:
+                choice_count = choice(range(8))
+            resources_id = sample(choices, choice_count)
+            for rid in resources_id:
+                data = {'assetpermission_id': perm.id}
+                data[relation_name] = rid
+                relations.append(through(**data))
+        through.objects.bulk_create(relations, ignore_conflicts=True)
+
+    def do_generate(self, batch, batch_size):
+        perms = []
+
+        for i in batch:
+            name = forgery_py.basic.text()
+            name = f'AssetPermission: {name}'
+            perm = AssetPermission(name=name, org_id=self.org.id)
+            perms.append(perm)
+        created = AssetPermission.objects.bulk_create(perms, ignore_conflicts=True)
+        self.set_users(created)
+        self.set_user_groups(created)
+        self.set_assets(created)
+        self.set_nodes(created)
+        self.set_system_users(created)

utils/generate_fake_data/resources/users.py

@@ -0,0 +1,62 @@
+from random import choice, sample
+import forgery_py
+
+from .base import FakeDataGenerator
+
+from users.models import *
+from orgs.models import OrganizationMember
+
+
+class UserGroupGenerator(FakeDataGenerator):
+    resource = 'usergroup'
+
+    def do_generate(self, batch, batch_size):
+        groups = []
+        for i in batch:
+            group_name = forgery_py.name.job_title()
+            groups.append(UserGroup(name=group_name, org_id=self.org.id))
+        UserGroup.objects.bulk_create(groups, ignore_conflicts=True)
+
+
+class UserGenerator(FakeDataGenerator):
+    resource = 'user'
+    roles: list
+    groups_id: list
+
+    def pre_generate(self):
+        self.roles = list(dict(User.ROLE.choices).keys())
+        self.groups_id = list(UserGroup.objects.all().values_list('id', flat=True))
+
+    def set_org(self, users):
+        relations = []
+        for u in users:
+            relations.append(OrganizationMember(
+                org_id=self.org.id,
+                user_id=u.id,
+            ))
+        OrganizationMember.objects.bulk_create(relations, ignore_conflicts=True)
+
+    def set_groups(self, users):
+        relations = []
+        for i in users:
+            groups_to_join = sample(self.groups_id, 3)
+            _relations = [User.groups.through(user_id=i.id, usergroup_id=gid) for gid in groups_to_join]
+            relations.extend(_relations)
+        User.groups.through.objects.bulk_create(relations, ignore_conflicts=True)
+
+    def do_generate(self, batch, batch_size):
+        users = []
+        for i in batch:
+            username = forgery_py.internet.user_name(True)
+            email = forgery_py.internet.email_address()
+            u = User(
+                username=username,
+                email=email,
+                name=username.title(),
+                role=choice(self.roles),
+                created_by='Faker'
+            )
+            users.append(u)
+        users = User.objects.bulk_create(users, ignore_conflicts=True)
+        self.set_org(users)
+        self.set_groups(users)