diff --git a/idhub/management/commands/create_example_data.py b/idhub/management/commands/create_example_data.py new file mode 100644 index 0000000..7b3f45c --- /dev/null +++ b/idhub/management/commands/create_example_data.py @@ -0,0 +1,195 @@ +import random +import string + +from django.core.management.base import BaseCommand +from django.db import IntegrityError +from idhub.models import Event, Rol, Service, UserRol +from idhub_auth.models import User + +DEFAULT_OBJECTS_CREATED = 30 +RANDOM_STRING_LENGTH = 30 +EMAIL_RANDOM_STRING_LENGTH = 10 + + +class Command(BaseCommand): + help = """ + Populate the database with dummy values. + You can either specify which model to create objects for, + or create objects for all models. + If no data is specified it will create 30 events, users, + services, roles, and user roles + + usage: create_example_data [-h] [--option args] + options: + --amount N Create N objects (includes events, users, services, + roles, and user roles) + --events N Create N events + --users N Create N users + --superusers N Create N superusers + --services N Create N services + --roles N Create N roles + --userroles N Create N user roles + --userrole u s Create a user role for user u and service s + """ + created_users = [] + created_roles = [] + created_services = [] + + def handle(self, *args, **options): + any_option_used = False + + if options["events"]: + self.create_events(options["events"]) + any_option_used = True + if options["users"]: + self.create_users(options["users"]) + any_option_used = True + if options["superusers"]: + self.create_superusers(options["superusers"]) + any_option_used = True + if options["services"]: + self.create_services(options["services"]) + any_option_used = True + if options["roles"]: + self.create_roles(options["roles"]) + any_option_used = True + if options["userroles"]: + self.create_user_roles(options["userroles"]) + any_option_used = True + if options["userrole"]: + user = options["userrole"][0] + service = options["userrole"][1] + self.create_user_roles(1, user, service) + any_option_used = True + + if options["amount"]: + self.create_all(options["amount"]) + any_option_used = True + + if not any_option_used: + self.create_all() + + def add_arguments(self, parser): + parser.add_argument("--amount", type=int, action='store') + parser.add_argument("--events", type=int) + parser.add_argument("--users", type=int) + parser.add_argument("--superusers", type=int) + parser.add_argument("--services", type=int) + parser.add_argument("--roles", type=int) + parser.add_argument("--userroles", type=int) + parser.add_argument("--userrole", nargs=2, type=str) + + def create_all(self, amount=DEFAULT_OBJECTS_CREATED): + self.create_events(amount) + self.create_users(amount) + self.create_roles(amount) + self.create_services(amount) + self.create_user_roles(amount) + + def create_events(self, amount, user=None): + created_event_amount = 0 + for value in range(0, amount): + try: + Event.objects.create( + type=random.randint(1, 30), + message=create_random_string(), + user=user + ) + created_event_amount += 1 + except IntegrityError: + self.stdout.write("Couldn't create event") + self.stdout.write(f"Created {created_event_amount} events") + + def create_users(self, amount): + created_user_amount = 0 + for value in range(0, amount): + email = create_random_string(EMAIL_RANDOM_STRING_LENGTH) + "@example.org" + try: + User.objects.create( + email=email, + # Could be improved, maybe using Faker + first_name=create_random_string(random.randint(5, 10)), + last_name=create_random_string(random.randint(5, 10)) + ) + self.created_users.append(email) + created_user_amount += 1 + except IntegrityError: + self.stdout.write("Couldn't create user " + email) + + self.stdout.write(f"Created {created_user_amount} users") + + def create_superusers(self, amount=0): + """Superusers can only be created from the specific command""" + created_superuser_amount = 0 + for value in range(0, amount): + email = create_random_string(EMAIL_RANDOM_STRING_LENGTH) + try: + User.objects.create_superuser(email) + created_superuser_amount += 1 + except IntegrityError: + self.stdout.write("Couldn't create superuser " + email) + self.stdout.write(f"Created {created_superuser_amount} users") + + def create_services(self, amount): + created_service_amount = 0 + for value in range(0, amount): + domain = create_random_string(random.randint(5, 15)) + try: + service = Service.objects.create( + domain=domain, + description=create_random_string( + random.randint(50, 100)) + ) + self.created_services.append(domain) + try: + associated_rol = Rol.objects.get(name=random.choice( + self.created_roles)) + service.rol.add(associated_rol.id) + except IntegrityError: + self.stdout.write( + f"Couldn't associate role with service {domain}") + created_service_amount += 1 + except IntegrityError: + self.stdout.write("Couldn't create service " + domain) + self.stdout.write(f"Created {created_service_amount} services") + + def create_roles(self, amount): + created_role_amount = 0 + for value in range(0, amount): + # Could be improved, maybe using Faker + name = create_random_string(random.randint(5, 10)) + try: + Rol.objects.create(name=name, + description=create_random_string( + random.randint(50, 100))) + created_role_amount += 1 + except IntegrityError: + self.stdout.write("Couldn't create role " + name) + self.created_roles.append(name) + self.stdout.write(f"Created {created_role_amount} roles") + + def create_user_roles(self, amount, user_id=None, service_id=None): + created_user_role_amount = 0 + user_id = user_id if user_id is not None else random.choice( + self.created_users) + service_id = service_id if service_id is not None else random.choice( + self.created_services) + for value in range(0, amount): + try: + user = User.objects.get(email=user_id) + service = Service.objects.get(domain=service_id) + UserRol.objects.create( + user=user, + service=service + ) + created_user_role_amount += 1 + except IntegrityError: + self.stdout.write("Couldn't create user role for user " + user.email) + user_id = random.choice(self.created_users) + service_id = random.choice(self.created_services) + self.stdout.write(f"Created {created_user_role_amount} user roles") + + +def create_random_string(string_length=RANDOM_STRING_LENGTH): + return ''.join(random.choices(string.ascii_uppercase + string.digits, + k=string_length))