Produits Entr'ouvert » Hobo

from django.core.management.base import BaseCommand

from django_rbac.utils import get_role_model, get_ou_model

from django.contrib.auth import get_user_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)

    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'])

        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('Provisionning {} roles.'.format(roles.count()))

        engine.notify_roles(ous, roles, full=True)

    def provision_users(self, engine, ous, batch_size=512):

        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:

                engine.notify_users(ous, users)

                users = list(qs.filter(id__gt=users[-1].pk)[:batch_size])

        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('Provisionning {} normal users.'.format(normal_users.count()))

        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('Provisionning {} admin users.'.format(admin_users.count()))

        do_provision(admin_users)