Merge pull request #1730 from garyd203/cognitoidp-groups

Implement groups for cognito-idp

moto/cognitoidp/models.py

@@ -13,8 +13,7 @@ from jose import jws
 
 from moto.compat import OrderedDict
 from moto.core import BaseBackend, BaseModel
-from .exceptions import NotAuthorizedError, ResourceNotFoundError, UserNotFoundError
-
+from .exceptions import GroupExistsException, NotAuthorizedError, ResourceNotFoundError, UserNotFoundError
 
 UserStatus = {
     "FORCE_CHANGE_PASSWORD": "FORCE_CHANGE_PASSWORD",
@@ -68,6 +67,7 @@ class CognitoIdpUserPool(BaseModel):
 
         self.clients = OrderedDict()
         self.identity_providers = OrderedDict()
+        self.groups = OrderedDict()
         self.users = OrderedDict()
         self.refresh_tokens = {}
         self.access_tokens = {}
@@ -220,6 +220,33 @@ class CognitoIdpIdentityProvider(BaseModel):
         return identity_provider_json
 
 
+class CognitoIdpGroup(BaseModel):
+
+    def __init__(self, user_pool_id, group_name, description, role_arn, precedence):
+        self.user_pool_id = user_pool_id
+        self.group_name = group_name
+        self.description = description or ""
+        self.role_arn = role_arn
+        self.precedence = precedence
+        self.last_modified_date = datetime.datetime.now()
+        self.creation_date = self.last_modified_date
+
+        # Users who are members of this group.
+        # Note that these links are bidirectional.
+        self.users = set()
+
+    def to_json(self):
+        return {
+            "GroupName": self.group_name,
+            "UserPoolId": self.user_pool_id,
+            "Description": self.description,
+            "RoleArn": self.role_arn,
+            "Precedence": self.precedence,
+            "LastModifiedDate": time.mktime(self.last_modified_date.timetuple()),
+            "CreationDate": time.mktime(self.creation_date.timetuple()),
+        }
+
+
 class CognitoIdpUser(BaseModel):
 
     def __init__(self, user_pool_id, username, password, status, attributes):
@@ -233,6 +260,10 @@ class CognitoIdpUser(BaseModel):
         self.create_date = datetime.datetime.utcnow()
         self.last_modified_date = datetime.datetime.utcnow()
 
+        # Groups this user is a member of.
+        # Note that these links are bidirectional.
+        self.groups = set()
+
     def _base_json(self):
         return {
             "UserPoolId": self.user_pool_id,
@@ -405,6 +436,72 @@ class CognitoIdpBackend(BaseBackend):
 
         del user_pool.identity_providers[name]
 
+    # Group
+    def create_group(self, user_pool_id, group_name, description, role_arn, precedence):
+        user_pool = self.user_pools.get(user_pool_id)
+        if not user_pool:
+            raise ResourceNotFoundError(user_pool_id)
+
+        group = CognitoIdpGroup(user_pool_id, group_name, description, role_arn, precedence)
+        if group.group_name in user_pool.groups:
+            raise GroupExistsException("A group with the name already exists")
+        user_pool.groups[group.group_name] = group
+
+        return group
+
+    def get_group(self, user_pool_id, group_name):
+        user_pool = self.user_pools.get(user_pool_id)
+        if not user_pool:
+            raise ResourceNotFoundError(user_pool_id)
+
+        if group_name not in user_pool.groups:
+            raise ResourceNotFoundError(group_name)
+
+        return user_pool.groups[group_name]
+
+    def list_groups(self, user_pool_id):
+        user_pool = self.user_pools.get(user_pool_id)
+        if not user_pool:
+            raise ResourceNotFoundError(user_pool_id)
+
+        return user_pool.groups.values()
+
+    def delete_group(self, user_pool_id, group_name):
+        user_pool = self.user_pools.get(user_pool_id)
+        if not user_pool:
+            raise ResourceNotFoundError(user_pool_id)
+
+        if group_name not in user_pool.groups:
+            raise ResourceNotFoundError(group_name)
+
+        group = user_pool.groups[group_name]
+        for user in group.users:
+            user.groups.remove(group)
+
+        del user_pool.groups[group_name]
+
+    def admin_add_user_to_group(self, user_pool_id, group_name, username):
+        group = self.get_group(user_pool_id, group_name)
+        user = self.admin_get_user(user_pool_id, username)
+
+        group.users.add(user)
+        user.groups.add(group)
+
+    def list_users_in_group(self, user_pool_id, group_name):
+        group = self.get_group(user_pool_id, group_name)
+        return list(group.users)
+
+    def admin_list_groups_for_user(self, user_pool_id, username):
+        user = self.admin_get_user(user_pool_id, username)
+        return list(user.groups)
+
+    def admin_remove_user_from_group(self, user_pool_id, group_name, username):
+        group = self.get_group(user_pool_id, group_name)
+        user = self.admin_get_user(user_pool_id, username)
+
+        group.users.discard(user)
+        user.groups.discard(group)
+
     # User
     def admin_create_user(self, user_pool_id, username, temporary_password, attributes):
         user_pool = self.user_pools.get(user_pool_id)
@@ -449,6 +546,10 @@ class CognitoIdpBackend(BaseBackend):
         if username not in user_pool.users:
             raise UserNotFoundError(username)
 
+        user = user_pool.users[username]
+        for group in user.groups:
+            group.users.remove(user)
+
         del user_pool.users[username]
 
     def _log_user_in(self, user_pool, client, username):