Repository: getmoto/moto

Issue:
mock_cognitoidp idtoken is missing AWS oauth headers
I am attempting to test a function that accepts a Cognito id token, fetches the public jwk for a cognito userpool, uses that key to decode and verify the token and extract the claims. It's based on sample code from an AWS tutorial. It works in a real environment with genuine Cognito id tokens, but I haven't yet gotten my moto unittests to work.

I have attempted to incorporate the test_token_verification example in test_cognitoidp.

The idtoken I am generating with my mock_cognito userpool and decoding with the provided jwks-public.json only has two headers: 

`{'alg': 'RS256', 'typ': 'JWT'}

In order to verify the cognito id token and extract claims, the code from AWS expects a 'kid' header, which appears in the jwks-public.json but is not in the headers.

From the token verification function:

> headers = jwt.get_unverified_headers(token)
        kid = headers['kid']
        

I seem to be encountering an issue related to the open PR [https://github.com/spulec/moto/pull/1689](https://github.com/spulec/moto/pull/1689)

I've attempted to implement this fix in my own unittest, building my own payload and headers and signing them with the dummy private key from this repo, but didn't make a lot of headway.




Discussion:
Hi @Alladinsane, can you share what you have so far in your unit test?
Sure thing.  Thanks for offering to take a look. The app I'm testing is a swagger-codegen flask API written in python3.7, using moto 1.3.14

test_Authorization_impl.py
```
from __future__ import absolute_import

import uuid

from moto import mock_cognitoidp

from swagger_server.test import BaseTestCase
from unittest.mock import patch
from swagger_server.test import test_environment


@mock_cognitoidp
class MyTestCase(BaseTestCase):

    def setUp(self):
        cognito = boto3.client('cognito-idp', 'us-west-2')
        self.user_pool_id = cognito.create_user_pool(PoolName=str(uuid.uuid4()))["UserPool"]["Id"]
        client = cognito.create_user_pool_client(
            UserPoolId=self.user_pool_id,
            ExplicitAuthFlows=[
                'ALLOW_ADMIN_USER_PASSWORD_AUTH',
                'ALLOW_REFRESH_TOKEN_AUTH',
                'ALLOW_ADMIN_NO_SRP_AUTH'
            ],
            AllowedOAuthFlows=['code', 'implicit'],
            ClientName=str(uuid.uuid4()),
            CallbackURLs=[
                'https://example.com'
            ]
        )
        self.client_id = client['UserPoolClient']['ClientId']
        self.username = str(uuid.uuid4())
        temporary_password = "1TemporaryP@ssword"
        self.cognito.admin_create_user(
            UserPoolId=self.user_pool_id,
            Username=self.username,
            TemporaryPassword=temporary_password
        )
        result = self.cognito.admin_initiate_auth(
            UserPoolId=self.user_pool_id,
            ClientId=self.client_id,
            AuthFlow="ADMIN_NO_SRP_AUTH",
            AuthParameters={"USERNAME": self.username, "PASSWORD": temporary_password}
        )

        # A newly created user is forced to set a new password
  
        # This sets a new password and logs the user in (creates tokens)
        self.password = "1F@kePassword"
        result = self.cognito.respond_to_auth_challenge(
            Session=result["Session"],
            ClientId=self.client_id,
            ChallengeName="NEW_PASSWORD_REQUIRED",
            ChallengeResponses={"USERNAME": self.username, "NEW_PASSWORD": self.password},
        )
         self.outputs = {
               "user_pool_id": self.user_pool_id,
               "client_id": self.client_id,
               "id_token": result["AuthenticationResult"]["IdToken"],
               "access_token": result["AuthenticationResult"]["AccessToken"],
               "username": self.username,
               "password": self.password,
         }

    def test_verifyToken(self):
        id_token = self.outputs['id_token']
        bad_token = 'jaoiealgadummytokenaoijgoain'

        from swagger_server.server_impl.controllers_impl import Authorization_impl
        with patch.multiple(Authorization_impl, COGNITO_CLIENT_ID=self.client_id,
                            USER_POOL=self.user_pool_id, cognito=self.cognito, 
                            REGION=self.region):
            with patch('swagger_server.server_impl.controllers_impl.constants.fetch_keys',
                       return_value=test_environment.fetch_keys()):
                response = Authorization_impl.verifyToken(id_token)
                print(response)

                self.assertIn('scope', response)

                response = Authorization_impl.verifyToken(bad_token)
                self.assert401(response)
`
```

Authorization_impl.verifyToken
```
def verifyToken(token):
    """Verifies token and returns a scope"""
        keys = constants.fetch_keys()
        headers = jwt.get_unverified_headers(token)
        kid = headers['kid']
        # At this point, the public key is checked for a matching kid header, 
        # but fails at this line with the moto resources because the id token
        # has no kid header

        key_index = -1
        for i in range(len(keys)):
            print(keys[i]['kid'])
            if kid == keys[i]['kid']:
                key_index = i
                break
        if key_index == -1:
            message = 'Public key not found in jwks.json'
            print(message)
            error = ErrorModel['code': 401, 'message': message]
            response = flask.make_response(
                flask.jsonify({'msg': error}), 401)
            return response
        # construct the public key
        public_key = jwk.construct(keys[key_index])
        # At this point, additional verification checks are made,
        # claims are decoded, and a scope is returned. 
        # This works with real cognito resources, both in the live app 
        # and with the unittest

```

constants.fetch_keys()
```
keys_url = 'https://cognito-idp.{}.amazonaws.com/{}/.well-known/jwks.json'\
    .format(REGION, USER_POOL)
print("url", keys_url)


def fetch_keys():
    with (urllib.request.urlopen(keys_url)) as f:
        response = f.read()
        my_keys = json.loads(response.decode('utf-8'))['keys']
        return my_keys
```

test_environment.fetch_keys()
```
def fetch_keys():
    path = "./jwks-public.json"
    with open(os.path.join(os.path.dirname(__file__), path)) as f:
        json_web_key = json.loads(f.read())["keys"]
    return json_web_key
```

When this failed, I took a stab at trying to use the moto dummy private key to sign my own idtoken, but had unknown errors with the decoding that I'll probably be able to eventually meddle through. However, I also have the expectation that I can probably configure moto to return me an idtoken with the appropriate headers, and have just overlooked something.

Resolving patch:
diff --git a/moto/cognitoidp/models.py b/moto/cognitoidp/models.py
--- a/moto/cognitoidp/models.py
+++ b/moto/cognitoidp/models.py
@@ -128,8 +128,12 @@ def create_jwt(
             "exp": now + expires_in,
         }
         payload.update(extra_data)
+        headers = {"kid": "dummy"}  # KID as present in jwks-public.json
 
-        return jws.sign(payload, self.json_web_key, algorithm="RS256"), expires_in
+        return (
+            jws.sign(payload, self.json_web_key, headers, algorithm="RS256"),
+            expires_in,
+        )
 
     def create_id_token(self, client_id, username):
         extra_data = self.get_user_extra_data_by_client_id(client_id, username)
diff --git a/moto/cognitoidp/urls.py b/moto/cognitoidp/urls.py
--- a/moto/cognitoidp/urls.py
+++ b/moto/cognitoidp/urls.py
@@ -5,5 +5,5 @@
 
 url_paths = {
     "{0}/$": CognitoIdpResponse.dispatch,
-    "{0}/<user_pool_id>/.well-known/jwks.json$": CognitoIdpJsonWebKeyResponse().serve_json_web_key,
+    "{0}/(?P<user_pool_id>[^/]+)/.well-known/jwks.json$": CognitoIdpJsonWebKeyResponse().serve_json_web_key,
 }


Test patch:
diff --git a/tests/test_cognitoidp/test_cognitoidp.py b/tests/test_cognitoidp/test_cognitoidp.py
--- a/tests/test_cognitoidp/test_cognitoidp.py
+++ b/tests/test_cognitoidp/test_cognitoidp.py
@@ -3,6 +3,7 @@
 import json
 import os
 import random
+import requests
 import uuid
 
 import boto3
@@ -10,10 +11,10 @@
 # noinspection PyUnresolvedReferences
 import sure  # noqa
 from botocore.exceptions import ClientError
-from jose import jws
+from jose import jws, jwk, jwt
 from nose.tools import assert_raises
 
-from moto import mock_cognitoidp
+from moto import mock_cognitoidp, settings
 from moto.core import ACCOUNT_ID
 
 
@@ -1309,3 +1310,80 @@ def test_admin_update_user_attributes():
             val.should.equal("Doe")
         elif attr["Name"] == "given_name":
             val.should.equal("Jane")
+
+
+# Test will retrieve public key from cognito.amazonaws.com/.well-known/jwks.json,
+# which isnt mocked in ServerMode
+if not settings.TEST_SERVER_MODE:
+
+    @mock_cognitoidp
+    def test_idtoken_contains_kid_header():
+        # https://github.com/spulec/moto/issues/3078
+        # Setup
+        cognito = boto3.client("cognito-idp", "us-west-2")
+        user_pool_id = cognito.create_user_pool(PoolName=str(uuid.uuid4()))["UserPool"][
+            "Id"
+        ]
+        client = cognito.create_user_pool_client(
+            UserPoolId=user_pool_id,
+            ExplicitAuthFlows=[
+                "ALLOW_ADMIN_USER_PASSWORD_AUTH",
+                "ALLOW_REFRESH_TOKEN_AUTH",
+                "ALLOW_ADMIN_NO_SRP_AUTH",
+            ],
+            AllowedOAuthFlows=["code", "implicit"],
+            ClientName=str(uuid.uuid4()),
+            CallbackURLs=["https://example.com"],
+        )
+        client_id = client["UserPoolClient"]["ClientId"]
+        username = str(uuid.uuid4())
+        temporary_password = "1TemporaryP@ssword"
+        cognito.admin_create_user(
+            UserPoolId=user_pool_id,
+            Username=username,
+            TemporaryPassword=temporary_password,
+        )
+        result = cognito.admin_initiate_auth(
+            UserPoolId=user_pool_id,
+            ClientId=client_id,
+            AuthFlow="ADMIN_NO_SRP_AUTH",
+            AuthParameters={"USERNAME": username, "PASSWORD": temporary_password},
+        )
+
+        # A newly created user is forced to set a new password
+        # This sets a new password and logs the user in (creates tokens)
+        password = "1F@kePassword"
+        result = cognito.respond_to_auth_challenge(
+            Session=result["Session"],
+            ClientId=client_id,
+            ChallengeName="NEW_PASSWORD_REQUIRED",
+            ChallengeResponses={"USERNAME": username, "NEW_PASSWORD": password},
+        )
+        #
+        id_token = result["AuthenticationResult"]["IdToken"]
+
+        # Verify the KID header is present in the token, and corresponds to the KID supplied by the public JWT
+        verify_kid_header(id_token)
+
+
+def verify_kid_header(token):
+    """Verifies the kid-header is corresponds with the public key"""
+    headers = jwt.get_unverified_headers(token)
+    kid = headers["kid"]
+
+    key_index = -1
+    keys = fetch_public_keys()
+    for i in range(len(keys)):
+        if kid == keys[i]["kid"]:
+            key_index = i
+            break
+    if key_index == -1:
+        raise Exception("Public key (kid) not found in jwks.json")
+
+
+def fetch_public_keys():
+    keys_url = "https://cognito-idp.{}.amazonaws.com/{}/.well-known/jwks.json".format(
+        "us-west-2", "someuserpoolid"
+    )
+    response = requests.get(keys_url).json()
+    return response["keys"]