77 lines
3.0 KiB
Python
77 lines
3.0 KiB
Python
import time
|
|
|
|
from django.contrib.auth import get_user_model
|
|
from django.core.management.base import BaseCommand
|
|
from django_rbac.utils import get_ou_model, get_role_model
|
|
|
|
from hobo.agent.authentic2.provisionning import Provisionning
|
|
|
|
|
|
class Command(BaseCommand):
|
|
help = 'Provision all roles or users'
|
|
|
|
def add_arguments(self, parser):
|
|
parser.add_argument('--roles', action='store_true', default=False)
|
|
parser.add_argument('--users', action='store_true', default=False)
|
|
parser.add_argument('--batch-size', type=int, default=512)
|
|
parser.add_argument('--batch-sleep', type=int, default=30)
|
|
|
|
def handle(self, *args, **options):
|
|
self.verbosity = options['verbosity']
|
|
engine = Provisionning()
|
|
ous = {ou.id: ou for ou in get_ou_model().objects.all()}
|
|
|
|
if options['roles']:
|
|
self.provision_roles(engine, ous)
|
|
|
|
if options['users']:
|
|
self.provision_users(
|
|
engine,
|
|
ous,
|
|
batch_size=options['batch_size'],
|
|
batch_sleep=options['batch_sleep'],
|
|
verbosity=options['verbosity'],
|
|
)
|
|
if self.verbosity > 0:
|
|
self.stdout.write('Done.')
|
|
|
|
def provision_roles(self, engine, ous):
|
|
roles = get_role_model().objects.all()
|
|
if self.verbosity > 0:
|
|
self.stdout.write(f'Provisionning {roles.count()} roles.')
|
|
engine.notify_roles(ous, roles, full=True)
|
|
|
|
def provision_users(self, engine, ous, batch_size=512, batch_sleep=30, verbosity=1):
|
|
qs = get_user_model().objects.all()
|
|
# allow easy pagination by pk
|
|
qs = qs.order_by('pk')
|
|
# prevent too much select
|
|
qs = qs.prefetch_related('attribute_values__attribute')
|
|
|
|
def do_provision(qs):
|
|
users = list(qs[:batch_size])
|
|
while users:
|
|
if verbosity > 0:
|
|
self.stdout.write(
|
|
' batch provisionning %d users and sleeping for %d seconds'
|
|
% (len(users), batch_sleep)
|
|
)
|
|
engine.notify_users(ous, users)
|
|
users = list(qs.filter(id__gt=users[-1].pk)[:batch_size])
|
|
if users:
|
|
time.sleep(batch_sleep)
|
|
|
|
roles_with_attributes = get_role_model().objects.filter(attributes__name='is_superuser').children()
|
|
# first those without and admin attribute
|
|
normal_users = qs.exclude(roles__in=roles_with_attributes)
|
|
|
|
if self.verbosity > 0:
|
|
self.stdout.write(f'Provisionning {normal_users.count()} normal users.')
|
|
do_provision(normal_users)
|
|
# then those with an admin attribute, use distinct to prevent
|
|
# duplicates caused by join on a m2m relation
|
|
admin_users = qs.filter(roles__in=roles_with_attributes).distinct()
|
|
if self.verbosity > 0:
|
|
self.stdout.write(f'Provisionning {admin_users.count()} admin users.')
|
|
do_provision(admin_users)
|