mirror of https://github.com/jumpserver/jumpserver
fit2bot
4 years ago
committed by
GitHub
17 changed files with 345 additions and 223 deletions
@ -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() |
||||
@ -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() |
||||
@ -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() |
||||
@ -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) |
||||
@ -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) |
||||
Loading…
Reference in new issue