jumpserver/utils/generate_fake_data/resources/perms.py

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'
    user_ids: list
    user_group_ids: list
    asset_ids: list
    node_ids: list
    system_user_ids: list

    def pre_generate(self):
        self.node_ids = list(Node.objects.all().values_list('id', flat=True))
        self.asset_ids = list(Asset.objects.all().values_list('id', flat=True))
        self.system_user_ids = list(SystemUser.objects.all().values_list('id', flat=True))
        self.user_ids = list(User.objects.all().values_list('id', flat=True))
        self.user_group_ids = list(UserGroup.objects.all().values_list('id', flat=True))

    def set_users(self, perms):
        through = AssetPermission.users.through
        choices = self.user_ids
        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_group_ids
        relation_name = 'usergroup_id'
        self.set_relations(perms, through, relation_name, choices)

    def set_assets(self, perms):
        through = AssetPermission.assets.through
        choices = self.asset_ids
        relation_name = 'asset_id'
        self.set_relations(perms, through, relation_name, choices)

    def set_nodes(self, perms):
        through = AssetPermission.nodes.through
        choices = self.node_ids
        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_user_ids
        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))
            resource_ids = sample(choices, choice_count)
            for rid in resource_ids:
                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)