Produits Entr'ouvert » Hobo

import time

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)

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

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

@patch('hobo.agent.authentic2.provisionning.notify_agents')

def test_command_hobo_provision(notify_agents, transactional_db, tenant, caplog):

    User = get_user_model()

    with tenant_context(tenant):

        ou = get_default_ou()

        LibertyProvider.objects.create(ou=ou, name='provider',

                                       entity_id='http://provider.com',

                                       protocol_conformance=lasso.PROTOCOL_SAML_2_0)

        for i in range(10):

            Role.objects.create(name='role-%s' % i, ou=ou)

        for i in range(10):

            User.objects.create(username='user-%s' % i, first_name='John',

                                last_name='Doe %s' % i, ou=ou,

                                email='jone.doe-%s@example.com')

    call_command('tenant_command', 'hobo_provision', domain=tenant.domain_url, roles=True, users=True)

    msg_1 = notify_agents.call_args_list[0][0][0]

    msg_2 = notify_agents.call_args_list[1][0][0]

    assert msg_1['@type'] == 'provision'

    assert msg_1['full'] is True

    assert msg_1['objects']['@type'] == 'role'

    assert len(msg_1['objects']['data']) == 10

    assert msg_2['@type'] == 'provision'

    assert msg_2['full'] is False

    assert msg_2['objects']['@type'] == 'user'

    assert len(msg_2['objects']['data']) == 10