From a91f8cf4b0ee3a756faccf1983f87c36a8522e21 Mon Sep 17 00:00:00 2001 From: Benjamin Dauvergne Date: Fri, 13 Apr 2018 16:30:12 +0200 Subject: [PATCH] agent/authentic2: add an hobo_provision command (#19853) To provision all users or roles, the role provisionning is full, i.e. it removes old roles. The user provisionning is not currently. --- .../management/commands/hobo_provision.py | 68 +++++++++++++++++++ tests_authentic/test_provisionning.py | 29 ++++++++ 2 files changed, 97 insertions(+) create mode 100644 hobo/agent/authentic2/management/commands/hobo_provision.py diff --git a/hobo/agent/authentic2/management/commands/hobo_provision.py b/hobo/agent/authentic2/management/commands/hobo_provision.py new file mode 100644 index 0000000..4ff456c --- /dev/null +++ b/hobo/agent/authentic2/management/commands/hobo_provision.py @@ -0,0 +1,68 @@ +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) diff --git a/tests_authentic/test_provisionning.py b/tests_authentic/test_provisionning.py index dab3799..62a44f9 100644 --- a/tests_authentic/test_provisionning.py +++ b/tests_authentic/test_provisionning.py @@ -453,3 +453,32 @@ def test_provision_createsuperuser(transactional_db, tenant, caplog): call_command('tenant_command', 'createsuperuser', domain=tenant.domain_url, uuid='coin', username='coin', email='coin@coin.org', interactive=False) assert notify_agents.call_count == 1 + + +@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 -- 2.17.0