IdHub/idhub/management/commands/create_example_data.py

199 lines
7.9 KiB
Python

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):
"""
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.
"""
help = 'Populate the database with dummy values for testing.'
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',
help='Create N objects (includes events, users, services, roles, and user roles)'
)
parser.add_argument('--events', type=int,
help='Create the specified number of events')
parser.add_argument('--users', type=int,
help='Create the specified number of users')
parser.add_argument('--superusers', type=int,
help='Create the specified number of superusers')
parser.add_argument('--services', type=int,
help='Create the specified number of services')
parser.add_argument('--roles', type=int,
help='Create the specified number of roles')
parser.add_argument('--userroles', type=int,
help='Create the specified number of user roles')
parser.add_argument(
'--userrole', nargs=2, type=str,
help='Create a user role for user U and service S',
metavar=('U', 'S'),
)
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))