|
1 |
from django.core.management.base import BaseCommand
|
|
2 |
|
|
3 |
from django_rbac.utils import get_role_model, get_ou_model
|
|
4 |
from django.contrib.auth import get_user_model
|
|
5 |
|
|
6 |
from hobo.agent.authentic2.provisionning import Provisionning
|
|
7 |
|
|
8 |
|
|
9 |
class Command(BaseCommand):
|
|
10 |
help = 'Provision all roles or users'
|
|
11 |
|
|
12 |
def add_arguments(self, parser):
|
|
13 |
parser.add_argument('--roles', action='store_true', default=False)
|
|
14 |
parser.add_argument('--users', action='store_true', default=False)
|
|
15 |
parser.add_argument('--batch-size', type=int, default=512)
|
|
16 |
|
|
17 |
def handle(self, *args, **options):
|
|
18 |
engine = Provisionning()
|
|
19 |
ous = {ou.id: ou for ou in get_ou_model().objects.all()}
|
|
20 |
|
|
21 |
if options['roles']:
|
|
22 |
self.provision_roles(engine, ous)
|
|
23 |
|
|
24 |
if options['users']:
|
|
25 |
self.provision_users(engine, ous, batch_size=options['batch_size'])
|
|
26 |
print 'Done.'
|
|
27 |
|
|
28 |
def provision_roles(self, engine, ous):
|
|
29 |
roles = get_role_model().objects.all()
|
|
30 |
print 'Provisionning', roles.count(), 'roles.'
|
|
31 |
engine.notify_roles(ous, roles, full=True)
|
|
32 |
|
|
33 |
def provision_users(self, engine, ous, batch_size=512):
|
|
34 |
qs = get_user_model().objects.all()
|
|
35 |
# allow easy pagination by pk
|
|
36 |
qs = qs.order_by('pk')
|
|
37 |
# prevent too much select
|
|
38 |
qs = qs.prefetch_related('attribute_values__attribute')
|
|
39 |
|
|
40 |
def do_provision(qs):
|
|
41 |
users = list(qs[:batch_size])
|
|
42 |
while users:
|
|
43 |
engine.notify_users(ous, users)
|
|
44 |
users = list(qs.filter(id__gt=users[-1].pk)[:batch_size])
|
|
45 |
|
|
46 |
roles_with_attributes = get_role_model().objects.filter(attributes__name='is_superuser').children()
|
|
47 |
# first those without and admin attribute
|
|
48 |
normal_users = qs.exclude(roles__in=roles_with_attributes)
|
|
49 |
print 'Provisionning', normal_users.count(), 'normal users.'
|
|
50 |
do_provision(normal_users)
|
|
51 |
# then thos with an admin attribute
|
|
52 |
admin_users = qs.filter(roles__in=roles_with_attributes)
|
|
53 |
print 'Provisionning', admin_users.count(), 'admin users.'
|
|
54 |
do_provision(admin_users)
|
0 |
|
-
|