From 87afbc52f6c7d5eb63f0da745625d05b01e81783 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sybren=20A=2E=20St=C3=BCvel?= <sybren@stuvel.eu>
Date: Tue, 22 Aug 2017 16:37:47 +0200
Subject: [PATCH] Updated do_badger to take an optional set of roles.

The 'role' parameter now must be passed as keyword arg instead of
positional arg. Either 'role' or 'roles' must be given.
---
 pillar/api/blender_cloud/subscription.py |  4 +--
 pillar/api/service.py                    | 43 +++++++++++++++++-------
 pillar/cli/setup.py                      |  2 +-
 3 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/pillar/api/blender_cloud/subscription.py b/pillar/api/blender_cloud/subscription.py
index 9d0eb7ab..63edbe47 100644
--- a/pillar/api/blender_cloud/subscription.py
+++ b/pillar/api/blender_cloud/subscription.py
@@ -118,7 +118,7 @@ def update_subscription():
         action = 'grant' if grant_subscriber else 'revoke'
         my_log.info('%sing subscriber role to user %s (Blender ID email %s)',
                     action, user_id, email)
-        service.do_badger(action, 'subscriber', user_id=user_id)
+        service.do_badger(action, role='subscriber', user_id=user_id)
     else:
         my_log.debug('Not changing subscriber role, grant=%r and is=%s',
                      grant_subscriber, is_subscriber)
@@ -126,7 +126,7 @@ def update_subscription():
     if grant_demo != is_demo:
         action = 'grant' if grant_demo else 'revoke'
         my_log.info('%sing demo role to user %s (Blender ID email %s)', action, user_id, email)
-        service.do_badger(action, 'demo', user_id=user_id)
+        service.do_badger(action, role='demo', user_id=user_id)
     else:
         my_log.debug('Not changing demo role, grant=%r and is=%s',
                      grant_demo, is_demo)
diff --git a/pillar/api/service.py b/pillar/api/service.py
index 1637a113..e97d043c 100644
--- a/pillar/api/service.py
+++ b/pillar/api/service.py
@@ -79,12 +79,15 @@ def badger():
                     action, user_email, role, action, role)
         return 'Role not allowed', 403
 
-    return do_badger(action, role, user_email=user_email)
+    return do_badger(action, role=role, user_email=user_email)
 
 
-def do_badger(action: str, role: str, *, user_email: str = '', user_id: bson.ObjectId = None):
+def do_badger(action: str, *,
+              role: str=None, roles: typing.Iterable[str]=None,
+              user_email: str = '', user_id: bson.ObjectId = None):
     """Performs a badger action, returning a HTTP response.
 
+    Either role or roles must be given.
     Either user_email or user_id must be given.
     """
 
@@ -98,9 +101,16 @@ def do_badger(action: str, role: str, *, user_email: str = '', user_id: bson.Obj
                   action, role, user_email, user_id)
         raise wz_exceptions.BadRequest('User email not given')
 
-    if not role:
-        log.error('do_badger(%r, %r, %r, %r): no role given.', action, role, user_email, user_id)
-        raise wz_exceptions.BadRequest('Role not given')
+    if bool(role) == bool(roles):
+        log.error('do_badger(%r, role=%r, roles=%r, %r, %r): '
+                  'either "role" or "roles" must be given.',
+                  action, role, roles, user_email, user_id)
+        raise wz_exceptions.BadRequest('Invalid role(s) given')
+
+    # If only a single role was given, handle it as a set of one role.
+    if not roles:
+        roles = {role}
+    del role
 
     users_coll = current_app.data.driver.db['users']
 
@@ -111,20 +121,29 @@ def do_badger(action: str, role: str, *, user_email: str = '', user_id: bson.Obj
         query = user_id
     db_user = users_coll.find_one(query, projection={'roles': 1, 'groups': 1})
     if db_user is None:
-        log.warning('badger(%s, %s, user_email=%s, user_id=%s): user not found',
-                    action, role, user_email, user_id)
+        log.warning('badger(%s, roles=%s, user_email=%s, user_id=%s): user not found',
+                    action, roles, user_email, user_id)
         return 'User not found', 404
 
     # Apply the action
-    roles = set(db_user.get('roles') or [])
+    user_roles = set(db_user.get('roles') or [])
     if action == 'grant':
-        roles.add(role)
+        user_roles |= roles
     else:
-        roles.discard(role)
+        user_roles -= roles
 
-    groups = manage_user_group_membership(db_user, role, action)
+    groups = None
+    for role in roles:
+        groups = manage_user_group_membership(db_user, role, action)
 
-    updates = {'roles': list(roles)}
+        if groups is None:
+            # No change for this role
+            continue
+
+        # Also update db_user for the next iteration.
+        db_user['groups'] = groups
+
+    updates = {'roles': list(user_roles)}
     if groups is not None:
         updates['groups'] = list(groups)
 
diff --git a/pillar/cli/setup.py b/pillar/cli/setup.py
index 539a2667..f553c937 100644
--- a/pillar/cli/setup.py
+++ b/pillar/cli/setup.py
@@ -85,7 +85,7 @@ def badger(action, user_email, role):
 
     with current_app.app_context():
         service.fetch_role_to_group_id_map()
-        response, status = service.do_badger(action, role, user_email=user_email)
+        response, status = service.do_badger(action, role=role, user_email=user_email)
 
     if status == 204:
         log.info('Done.')