Implement some of cognito-idp

moto/cognitoidp/models.py

@@ -0,0 +1,512 @@
+from __future__ import unicode_literals
+
+import datetime
+import json
+import os
+import time
+import uuid
+
+import boto.cognito.identity
+from jose import jws
+
+from moto.compat import OrderedDict
+from moto.core import BaseBackend, BaseModel
+from .exceptions import NotAuthorizedError, ResourceNotFoundError, UserNotFoundError
+
+
+UserStatus = {
+    "FORCE_CHANGE_PASSWORD": "FORCE_CHANGE_PASSWORD",
+    "CONFIRMED": "CONFIRMED",
+}
+
+
+class CognitoIdpUserPool(BaseModel):
+
+    def __init__(self, region, name, extended_config):
+        self.region = region
+        self.id = str(uuid.uuid4())
+        self.name = name
+        self.status = None
+        self.extended_config = extended_config or {}
+        self.creation_date = datetime.datetime.utcnow()
+        self.last_modified_date = datetime.datetime.utcnow()
+
+        self.clients = OrderedDict()
+        self.identity_providers = OrderedDict()
+        self.users = OrderedDict()
+        self.refresh_tokens = {}
+        self.access_tokens = {}
+        self.id_tokens = {}
+
+        with open(os.path.join(os.path.dirname(__file__), "resources/jwks-private.json")) as f:
+            self.json_web_key = json.loads(f.read())
+
+    def _base_json(self):
+        return {
+            "Id": self.id,
+            "Name": self.name,
+            "Status": self.status,
+            "CreationDate": time.mktime(self.creation_date.timetuple()),
+            "LastModifiedDate": time.mktime(self.last_modified_date.timetuple()),
+        }
+
+    def to_json(self, extended=False):
+        user_pool_json = self._base_json()
+        if extended:
+            user_pool_json.update(self.extended_config)
+        else:
+            user_pool_json["LambdaConfig"] = self.extended_config.get("LambdaConfig") or {}
+
+        return user_pool_json
+
+    def create_jwt(self, client_id, username, expires_in=60 * 60, extra_data={}):
+        now = int(time.time())
+        payload = {
+            "iss": "https://cognito-idp.{}.amazonaws.com/{}".format(self.region, self.id),
+            "sub": self.users[username].id,
+            "aud": client_id,
+            "token_use": "id",
+            "auth_time": now,
+            "exp": now + expires_in,
+        }
+        payload.update(extra_data)
+
+        return jws.sign(payload, self.json_web_key, algorithm='RS256'), expires_in
+
+    def create_id_token(self, client_id, username):
+        id_token, expires_in = self.create_jwt(client_id, username)
+        self.id_tokens[id_token] = (client_id, username)
+        return id_token, expires_in
+
+    def create_refresh_token(self, client_id, username):
+        refresh_token = str(uuid.uuid4())
+        self.refresh_tokens[refresh_token] = (client_id, username)
+        return refresh_token
+
+    def create_access_token(self, client_id, username):
+        access_token, expires_in = self.create_jwt(client_id, username)
+        self.access_tokens[access_token] = (client_id, username)
+        return access_token, expires_in
+
+    def create_tokens_from_refresh_token(self, refresh_token):
+        client_id, username = self.refresh_tokens.get(refresh_token)
+        if not username:
+            raise NotAuthorizedError(refresh_token)
+
+        access_token, expires_in = self.create_access_token(client_id, username)
+        id_token, _ = self.create_id_token(client_id, username)
+        return access_token, id_token, expires_in
+
+
+class CognitoIdpUserPoolDomain(BaseModel):
+
+    def __init__(self, user_pool_id, domain):
+        self.user_pool_id = user_pool_id
+        self.domain = domain
+
+    def to_json(self):
+        return {
+            "UserPoolId": self.user_pool_id,
+            "AWSAccountId": str(uuid.uuid4()),
+            "CloudFrontDistribution": None,
+            "Domain": self.domain,
+            "S3Bucket": None,
+            "Status": "ACTIVE",
+            "Version": None,
+        }
+
+
+class CognitoIdpUserPoolClient(BaseModel):
+
+    def __init__(self, user_pool_id, extended_config):
+        self.user_pool_id = user_pool_id
+        self.id = str(uuid.uuid4())
+        self.secret = str(uuid.uuid4())
+        self.extended_config = extended_config or {}
+
+    def _base_json(self):
+        return {
+            "ClientId": self.id,
+            "ClientName": self.extended_config.get("ClientName"),
+            "UserPoolId": self.user_pool_id,
+        }
+
+    def to_json(self, extended=False):
+        user_pool_client_json = self._base_json()
+        if extended:
+            user_pool_client_json.update(self.extended_config)
+
+        return user_pool_client_json
+
+
+class CognitoIdpIdentityProvider(BaseModel):
+
+    def __init__(self, name, extended_config):
+        self.name = name
+        self.extended_config = extended_config or {}
+        self.creation_date = datetime.datetime.utcnow()
+        self.last_modified_date = datetime.datetime.utcnow()
+
+    def _base_json(self):
+        return {
+            "ProviderName": self.name,
+            "ProviderType": self.extended_config.get("ProviderType"),
+            "CreationDate": time.mktime(self.creation_date.timetuple()),
+            "LastModifiedDate": time.mktime(self.last_modified_date.timetuple()),
+        }
+
+    def to_json(self, extended=False):
+        identity_provider_json = self._base_json()
+        if extended:
+            identity_provider_json.update(self.extended_config)
+
+        return identity_provider_json
+
+
+class CognitoIdpUser(BaseModel):
+
+    def __init__(self, user_pool_id, username, password, status, attributes):
+        self.id = str(uuid.uuid4())
+        self.user_pool_id = user_pool_id
+        self.username = username
+        self.password = password
+        self.status = status
+        self.enabled = True
+        self.attributes = attributes
+        self.create_date = datetime.datetime.utcnow()
+        self.last_modified_date = datetime.datetime.utcnow()
+
+    def _base_json(self):
+        return {
+            "UserPoolId": self.user_pool_id,
+            "Username": self.username,
+            "UserStatus": self.status,
+            "UserCreateDate": time.mktime(self.create_date.timetuple()),
+            "UserLastModifiedDate": time.mktime(self.last_modified_date.timetuple()),
+        }
+
+    # list_users brings back "Attributes" while admin_get_user brings back "UserAttributes".
+    def to_json(self, extended=False, attributes_key="Attributes"):
+        user_json = self._base_json()
+        if extended:
+            user_json.update(
+                {
+                    "Enabled": self.enabled,
+                    attributes_key: self.attributes,
+                    "MFAOptions": []
+                }
+            )
+
+        return user_json
+
+
+class CognitoIdpBackend(BaseBackend):
+
+    def __init__(self, region):
+        super(CognitoIdpBackend, self).__init__()
+        self.region = region
+        self.user_pools = OrderedDict()
+        self.user_pool_domains = OrderedDict()
+        self.sessions = {}
+
+    def reset(self):
+        region = self.region
+        self.__dict__ = {}
+        self.__init__(region)
+
+    # User pool
+    def create_user_pool(self, name, extended_config):
+        user_pool = CognitoIdpUserPool(self.region, name, extended_config)
+        self.user_pools[user_pool.id] = user_pool
+        return user_pool
+
+    def list_user_pools(self):
+        return self.user_pools.values()
+
+    def describe_user_pool(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
+
+    def delete_user_pool(self, user_pool_id):
+        if user_pool_id not in self.user_pools:
+            raise ResourceNotFoundError(user_pool_id)
+
+        del self.user_pools[user_pool_id]
+
+    # User pool domain
+    def create_user_pool_domain(self, user_pool_id, domain):
+        if user_pool_id not in self.user_pools:
+            raise ResourceNotFoundError(user_pool_id)
+
+        user_pool_domain = CognitoIdpUserPoolDomain(user_pool_id, domain)
+        self.user_pool_domains[domain] = user_pool_domain
+        return user_pool_domain
+
+    def describe_user_pool_domain(self, domain):
+        if domain not in self.user_pool_domains:
+            return None
+
+        return self.user_pool_domains[domain]
+
+    def delete_user_pool_domain(self, domain):
+        if domain not in self.user_pool_domains:
+            raise ResourceNotFoundError(domain)
+
+        del self.user_pool_domains[domain]
+
+    # User pool client
+    def create_user_pool_client(self, user_pool_id, extended_config):
+        user_pool = self.user_pools.get(user_pool_id)
+        if not user_pool:
+            raise ResourceNotFoundError(user_pool_id)
+
+        user_pool_client = CognitoIdpUserPoolClient(user_pool_id, extended_config)
+        user_pool.clients[user_pool_client.id] = user_pool_client
+        return user_pool_client
+
+    def list_user_pool_clients(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.clients.values()
+
+    def describe_user_pool_client(self, user_pool_id, client_id):
+        user_pool = self.user_pools.get(user_pool_id)
+        if not user_pool:
+            raise ResourceNotFoundError(user_pool_id)
+
+        client = user_pool.clients.get(client_id)
+        if not client:
+            raise ResourceNotFoundError(client_id)
+
+        return client
+
+    def update_user_pool_client(self, user_pool_id, client_id, extended_config):
+        user_pool = self.user_pools.get(user_pool_id)
+        if not user_pool:
+            raise ResourceNotFoundError(user_pool_id)
+
+        client = user_pool.clients.get(client_id)
+        if not client:
+            raise ResourceNotFoundError(client_id)
+
+        client.extended_config.update(extended_config)
+        return client
+
+    def delete_user_pool_client(self, user_pool_id, client_id):
+        user_pool = self.user_pools.get(user_pool_id)
+        if not user_pool:
+            raise ResourceNotFoundError(user_pool_id)
+
+        if client_id not in user_pool.clients:
+            raise ResourceNotFoundError(client_id)
+
+        del user_pool.clients[client_id]
+
+    # Identity provider
+    def create_identity_provider(self, user_pool_id, name, extended_config):
+        user_pool = self.user_pools.get(user_pool_id)
+        if not user_pool:
+            raise ResourceNotFoundError(user_pool_id)
+
+        identity_provider = CognitoIdpIdentityProvider(name, extended_config)
+        user_pool.identity_providers[name] = identity_provider
+        return identity_provider
+
+    def list_identity_providers(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.identity_providers.values()
+
+    def describe_identity_provider(self, user_pool_id, name):
+        user_pool = self.user_pools.get(user_pool_id)
+        if not user_pool:
+            raise ResourceNotFoundError(user_pool_id)
+
+        identity_provider = user_pool.identity_providers.get(name)
+        if not identity_provider:
+            raise ResourceNotFoundError(name)
+
+        return identity_provider
+
+    def delete_identity_provider(self, user_pool_id, name):
+        user_pool = self.user_pools.get(user_pool_id)
+        if not user_pool:
+            raise ResourceNotFoundError(user_pool_id)
+
+        if name not in user_pool.identity_providers:
+            raise ResourceNotFoundError(name)
+
+        del user_pool.identity_providers[name]
+
+    # User
+    def admin_create_user(self, user_pool_id, username, temporary_password, attributes):
+        user_pool = self.user_pools.get(user_pool_id)
+        if not user_pool:
+            raise ResourceNotFoundError(user_pool_id)
+
+        user = CognitoIdpUser(user_pool_id, username, temporary_password, UserStatus["FORCE_CHANGE_PASSWORD"], attributes)
+        user_pool.users[user.username] = user
+        return user
+
+    def admin_get_user(self, user_pool_id, username):
+        user_pool = self.user_pools.get(user_pool_id)
+        if not user_pool:
+            raise ResourceNotFoundError(user_pool_id)
+
+        if username not in user_pool.users:
+            raise ResourceNotFoundError(username)
+
+        return user_pool.users[username]
+
+    def list_users(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.users.values()
+
+    def admin_delete_user(self, user_pool_id, username):
+        user_pool = self.user_pools.get(user_pool_id)
+        if not user_pool:
+            raise ResourceNotFoundError(user_pool_id)
+
+        if username not in user_pool.users:
+            raise ResourceNotFoundError(username)
+
+        del user_pool.users[username]
+
+    def _log_user_in(self, user_pool, client, username):
+        refresh_token = user_pool.create_refresh_token(client.id, username)
+        access_token, id_token, expires_in = user_pool.create_tokens_from_refresh_token(refresh_token)
+
+        return {
+            "AuthenticationResult": {
+                "IdToken": id_token,
+                "AccessToken": access_token,
+                "RefreshToken": refresh_token,
+                "ExpiresIn": expires_in,
+            }
+        }
+
+    def admin_initiate_auth(self, user_pool_id, client_id, auth_flow, auth_parameters):
+        user_pool = self.user_pools.get(user_pool_id)
+        if not user_pool:
+            raise ResourceNotFoundError(user_pool_id)
+
+        client = user_pool.clients.get(client_id)
+        if not client:
+            raise ResourceNotFoundError(client_id)
+
+        if auth_flow == "ADMIN_NO_SRP_AUTH":
+            username = auth_parameters.get("USERNAME")
+            password = auth_parameters.get("PASSWORD")
+            user = user_pool.users.get(username)
+            if not user:
+                raise UserNotFoundError(username)
+
+            if user.password != password:
+                raise NotAuthorizedError(username)
+
+            if user.status == UserStatus["FORCE_CHANGE_PASSWORD"]:
+                session = str(uuid.uuid4())
+                self.sessions[session] = user_pool
+
+                return {
+                    "ChallengeName": "NEW_PASSWORD_REQUIRED",
+                    "ChallengeParameters": {},
+                    "Session": session,
+                }
+
+            return self._log_user_in(user_pool, client, username)
+        elif auth_flow == "REFRESH_TOKEN":
+            refresh_token = auth_parameters.get("REFRESH_TOKEN")
+            id_token, access_token, expires_in = user_pool.create_tokens_from_refresh_token(refresh_token)
+
+            return {
+                "AuthenticationResult": {
+                    "IdToken": id_token,
+                    "AccessToken": access_token,
+                    "ExpiresIn": expires_in,
+                }
+            }
+        else:
+            return {}
+
+    def respond_to_auth_challenge(self, session, client_id, challenge_name, challenge_responses):
+        user_pool = self.sessions.get(session)
+        if not user_pool:
+            raise ResourceNotFoundError(session)
+
+        client = user_pool.clients.get(client_id)
+        if not client:
+            raise ResourceNotFoundError(client_id)
+
+        if challenge_name == "NEW_PASSWORD_REQUIRED":
+            username = challenge_responses.get("USERNAME")
+            new_password = challenge_responses.get("NEW_PASSWORD")
+            user = user_pool.users.get(username)
+            if not user:
+                raise UserNotFoundError(username)
+
+            user.password = new_password
+            user.status = UserStatus["CONFIRMED"]
+            del self.sessions[session]
+
+            return self._log_user_in(user_pool, client, username)
+        else:
+            return {}
+
+    def confirm_forgot_password(self, client_id, username, password):
+        for user_pool in self.user_pools.values():
+            if client_id in user_pool.clients and username in user_pool.users:
+                user_pool.users[username].password = password
+                break
+        else:
+            raise ResourceNotFoundError(client_id)
+
+    def change_password(self, access_token, previous_password, proposed_password):
+        for user_pool in self.user_pools.values():
+            if access_token in user_pool.access_tokens:
+                _, username = user_pool.access_tokens[access_token]
+                user = user_pool.users.get(username)
+                if not user:
+                    raise UserNotFoundError(username)
+
+                if user.password != previous_password:
+                    raise NotAuthorizedError(username)
+
+                user.password = proposed_password
+                if user.status == UserStatus["FORCE_CHANGE_PASSWORD"]:
+                    user.status = UserStatus["CONFIRMED"]
+
+                break
+        else:
+            raise NotAuthorizedError(access_token)
+
+
+cognitoidp_backends = {}
+for region in boto.cognito.identity.regions():
+    cognitoidp_backends[region.name] = CognitoIdpBackend(region.name)
+
+
+# Hack to help moto-server process requests on localhost, where the region isn't
+# specified in the host header. Some endpoints (change password, confirm forgot
+# password) have no authorization header from which to extract the region.
+def find_region_by_value(key, value):
+    for region in cognitoidp_backends:
+        backend = cognitoidp_backends[region]
+        for user_pool in backend.user_pools.values():
+            if key == "client_id" and value in user_pool.clients:
+                return region
+
+            if key == "access_token" and value in user_pool.access_tokens:
+                return region
+
+    return cognitoidp_backends.keys()[0]