diff --git a/utils/generate_fake_data/generate.py b/utils/generate_fake_data/generate.py index fb41d7e1d..97340e8e6 100644 --- a/utils/generate_fake_data/generate.py +++ b/utils/generate_fake_data/generate.py @@ -15,8 +15,7 @@ django.setup() from resources.assets import AssetsGenerator, NodesGenerator, SystemUsersGenerator, AdminUsersGenerator from resources.users import UserGroupGenerator, UserGenerator from resources.perms import AssetPermissionGenerator -from resources.terminal import CommandGenerator -# from resources.system import StatGenerator +from resources.terminal import CommandGenerator, SessionGenerator resource_generator_mapper = { @@ -28,6 +27,7 @@ resource_generator_mapper = { 'user_group': UserGroupGenerator, 'asset_permission': AssetPermissionGenerator, 'command': CommandGenerator, + 'session': SessionGenerator # 'stat': StatGenerator } diff --git a/utils/generate_fake_data/resources/system.py b/utils/generate_fake_data/resources/system.py index 69a12d144..fb40ce83a 100644 --- a/utils/generate_fake_data/resources/system.py +++ b/utils/generate_fake_data/resources/system.py @@ -1,7 +1,7 @@ import random from .base import FakeDataGenerator -from system.models import * +from terminal.models import * class StatGenerator(FakeDataGenerator): @@ -66,4 +66,4 @@ class StatGenerator(FakeDataGenerator): 'datetime': datetime }) items.append(Stat(**node)) - Stat.objects.bulk_create(items, ignore_conflicts=True) \ No newline at end of file + Stat.objects.bulk_create(items, ignore_conflicts=True) diff --git a/utils/generate_fake_data/resources/terminal.py b/utils/generate_fake_data/resources/terminal.py index 48d825b47..481e98f6d 100644 --- a/utils/generate_fake_data/resources/terminal.py +++ b/utils/generate_fake_data/resources/terminal.py @@ -1,5 +1,9 @@ from .base import FakeDataGenerator -from terminal.models import Command +from users.models import * +from assets.models import * +from terminal.models import * +import forgery_py +import random class CommandGenerator(FakeDataGenerator): @@ -8,3 +12,42 @@ class CommandGenerator(FakeDataGenerator): def do_generate(self, batch, batch_size): Command.generate_fake(len(batch), self.org) + +class SessionGenerator(FakeDataGenerator): + resource = 'session' + users: list + assets: list + system_users: list + + def pre_generate(self): + self.users = list(User.objects.all()) + self.assets = list(Asset.objects.all()) + self.system_users = list(SystemUser.objects.all()) + + def do_generate(self, batch, batch_size): + user = random.choice(self.users) + asset = random.choice(self.assets) + system_user = random.choice(self.system_users) + + objects = [] + + now = timezone.now() + timedelta = list(range(30)) + for i in batch: + delta = random.choice(timedelta) + date_start = now - timezone.timedelta(days=delta, seconds=delta * 60) + date_end = date_start + timezone.timedelta(seconds=delta * 60) + data = dict( + user=user.name, + user_id=user.id, + asset=asset.hostname, + asset_id=asset.id, + system_user=system_user.name, + system_user_id=system_user.id, + org_id=self.org.id, + date_start=date_start, + date_end=date_end, + ) + objects.append(Session(**data)) + creates = Session.objects.bulk_create(objects, ignore_conflicts=True) + diff --git a/utils/generate_fake_data/resources/users.py b/utils/generate_fake_data/resources/users.py index 887c7240c..c973fec5e 100644 --- a/utils/generate_fake_data/resources/users.py +++ b/utils/generate_fake_data/resources/users.py @@ -24,7 +24,6 @@ class UserGenerator(FakeDataGenerator): group_ids: list def pre_generate(self): - self.roles = list(dict(User.ROLE.choices).keys()) self.group_ids = list(UserGroup.objects.all().values_list('id', flat=True)) def set_org(self, users): @@ -53,7 +52,6 @@ class UserGenerator(FakeDataGenerator): username=username, email=email, name=username.title(), - role=choice(self.roles), created_by='Faker' ) users.append(u)