Add ability to inject extra claims and xsrfToken

stormpath/api_auth.py

@@ -5,6 +5,7 @@
 import datetime
 import json
 from six import string_types
+import uuid
 
 try:
     from urlparse import urlparse, parse_qs
@@ -94,10 +95,11 @@ def __init__(self, app, token):
 
         # get raw data without validation
         try:
-            data = jwt.decode(self.token, verify=False, algorithms=['HS256'])
-            self.client_id = data.get('sub', '')
+            claims = jwt.decode(self.token, verify=False, algorithms=['HS256'])
+            self.claims = claims
+            self.client_id = claims.get('sub', '')
             try:
-                self.account = self.app.accounts.get(data.get('sub', ''))
+                self.account = self.app.accounts.get(claims.get('sub', ''))
 
                 # We're accessing account.username here to force
                 # evaluation of this Account -- this allows us to check
@@ -109,8 +111,8 @@ def __init__(self, app, token):
             if self.account:
                 self.for_api_key = False
             self.api_key = self.app.api_keys.get_key(self.client_id)
-            self.exp = data.get('exp', 0)
-            self.scopes = data.get('scope', '') if data.get('scope') else ''
+            self.exp = claims.get('exp', 0)
+            self.scopes = claims.get('scope', '') if claims.get('scope') else ''
             self.scopes = self.scopes.split(' ')
         except jwt.DecodeError:
             pass
@@ -294,47 +296,104 @@ def validate_bearer_token(self, token, scopes, request):
             return True
         return False
 
+class JWTFactory(object):
+    def __init__(self, stormpath_app, token_data=None, ttl=None, add_xsrf_token=False, encode_params=None):
+        self.stormpath_app = stormpath_app
+        self.token_data = token_data
+        self.ttl = ttl
+        self.add_xsrf_token = add_xsrf_token
+        self.encode_params = encode_params
+
+    def get_bearer_token(self, *args, **kwargs):
+        kwargs.setdefault("token_generator", self.generate_signed_token)
+        return self._get_bearer_token(*args, **kwargs)
+
+    def generate_signed_token(self, request, **kwargs):
+        kwargs.setdefault("app", self.stormpath_app)
+        kwargs.setdefault("token_data", self.token_data)
+        kwargs.setdefault("ttl", self.ttl)
+        kwargs.setdefault("add_xsrf_token", self.add_xsrf_token)
+        kwargs.setdefault("encode_params", self.encode_params)
+        return self._generate_signed_token(
+            request,
+            **kwargs
+        )
+
+    @classmethod
+    def _generate_signed_token(cls, request, app=None, token_data=None,
+                               ttl=None, add_xsrf_token=False,
+                               encode_params=None, secret=None):
+        client_id = request.client.client_id
+        if app is None:
+            app = request.app
+
+        app.api_keys.get_key(client_id)
+        # the SP ApiKey is already validated in SPOauth2RequestValidator.validate_client_id
+        # but to prevent time based attacks oauthlib always goes through the entire
+        # flow  even though the entire request will be deemed invalid
+        # in the end.
+        if secret is None:
+            secret = app._client.auth.secret
+
+        now = datetime.datetime.utcnow()
+
+        if ttl is None:
+            ttl = request.expires_in
+
+        exp = now + datetime.timedelta(seconds=ttl)
+
+        claims = {
+            'iss': request.app.href,
+            'sub': client_id,
+            'iat': now,
+            'exp': exp
+        }
 
-def _generate_signed_token(request):
-    client_id = request.client.client_id
-    request.app.api_keys.get_key(client_id)
+        if hasattr(request, 'scope'):
+            claims['scope'] = request.scope
 
-    # the SP ApiKey is already validated in SPOauth2RequestValidator.validate_client_id
-    # but to prevent time based attacks oauthlib always goes through the entire
-    # flow  even though the entire request will be deemed invalid
-    # in the end.
-    secret = request.app._client.auth.secret
+        if add_xsrf_token:
+            claims["xsrfToken"] = str(uuid.uuid4().hex)
 
-    now = datetime.datetime.utcnow()
+        if token_data:
+            # Don't allow token_data to overwrite built-in claims
+            token_data.update(claims)
+            claims = token_data
+        if encode_params is None:
+            encode_params = {"algorithm":"HS256"}
+        token = jwt.encode(claims, secret, **encode_params)
+        token = to_unicode(token, "UTF-8")
 
-    data = {
-        'iss': request.app.href,
-        'sub': client_id,
-        'iat': now,
-        'exp': now + datetime.timedelta(seconds=request.expires_in)
-    }
+        return token
 
-    if hasattr(request, 'scope'):
-        data['scope'] = request.scope
 
-    token = jwt.encode(data, secret, 'HS256')
-    token = to_unicode(token, "UTF-8")
+    @classmethod
+    def _get_bearer_token(cls, app, allowed_scopes, http_method, uri, body, headers, ttl=DEFAULT_TTL,
+                          token_generator=None):
 
-    return token
+        validator = SPOauth2RequestValidator(app=app, allowed_scopes=allowed_scopes, ttl=ttl)
 
+        if token_generator is None:
+            token_generator = cls._generate_signed_token
 
-def _get_bearer_token(app, allowed_scopes, http_method, uri, body, headers, ttl=DEFAULT_TTL):
-    validator = SPOauth2RequestValidator(app=app, allowed_scopes=allowed_scopes, ttl=ttl)
-    server = Oauth2BackendApplicationServer(validator,
-        token_generator=_generate_signed_token)
+        server = Oauth2BackendApplicationServer(validator,
+            token_generator=token_generator)
 
-    headers, body, status = server.create_token_response(
-        uri, http_method, body, headers, {})
+        headers, body, status = server.create_token_response(
+            uri, http_method, body, headers, {})
+
+        if status == 200:
+            token_response = json.loads(body)
+            return token_response.get('access_token')
+        return None
 
-    if status == 200:
-        token_response = json.loads(body)
-        return token_response.get('access_token')
-    return None
+##########################################################
+# Deprecated
+##########################################################
+##########################################################
+_generate_signed_token = JWTFactory._generate_signed_token
+_get_bearer_token = JWTFactory._get_bearer_token
+##########################################################
 
 
 class Authenticator(object):
@@ -343,6 +402,16 @@ class Authenticator(object):
     :param app: An application to which this Authenticator
         authenticates.
 
+    :param token_data dict: A dictionary of extra data to be
+        added to JWT's if your authenticator generates them.
+        This will be merged into the JWT claims.  Any reserved
+        keys like 'exp' or 'iat' will not be overwritten.
+
+    :param add_xsrf_token bool: If True, an JWT's generated
+        by this authenticator will include a 'xsrfToken' key.
+        For more information, see
+        https://stormpath.com/blog/where-to-store-your-jwts-cookies-vs-html5-web-storage/
+
     """
     def __init__(self, app):
         self.app = app
@@ -409,8 +478,13 @@ def authenticate(self, headers, http_method='', uri='', body=None,
             account=result.account, api_key=result.api_key,
             access_token=access_token)
 
+class JWTRequestAuthenticator(RequestAuthenticator):
+    def __init__(self, stormpath_app, token_data=None, add_xsrf_token=False):
+        super(JWTRequestAuthenticator, self).__init__(stormpath_app)
+        self.jwt_factory = JWTFactory(stormpath_app, token_data=token_data, add_xsrf_token=add_xsrf_token)
+
 
-class ApiRequestAuthenticator(RequestAuthenticator):
+class ApiRequestAuthenticator(JWTRequestAuthenticator):
     """This class should authenticate both HTTP Basic Auth and OAuth2
     requests. However, if you need more specific or customized OAuth2
     request processing, you will likely want to use the
@@ -429,7 +503,7 @@ def _get_scheme_and_token(self, headers, http_method, uri, body, scopes,
             jwt_token = None
             if auth_scheme == 'Basic':
                 if body.get('grant_type') or url_qs.get('grant_type'):
-                    jwt_token =_get_bearer_token(
+                    jwt_token =self.jwt_factory.get_bearer_token(
                         self.app, scopes, http_method, uri, body, headers, ttl)
             if auth_scheme == 'Bearer':
                 jwt_token = auth_header.split(' ')[1]
@@ -455,7 +529,7 @@ def _get_scheme_and_token(self, headers, http_method, uri, body, scopes,
         return None, None
 
 
-class OAuthRequestAuthenticator(RequestAuthenticator):
+class OAuthRequestAuthenticator(JWTRequestAuthenticator):
     """This class should authenticate OAuth2 requests. It will
     eventually support authenticating all 4 OAuth2 grant types.
     Specifically, right now, this class will authenticate OAuth2
@@ -476,7 +550,7 @@ def _get_scheme_and_token(self, headers, http_method, uri, body, scopes,
             if auth_scheme == 'Basic':
                 if not 'grant_type' in body and not 'grant_type' in url_qs:
                     return None, None
-                jwt_token =_get_bearer_token(
+                jwt_token =self.jwt_factory.get_bearer_token(
                     self.app, scopes, http_method, uri, body, headers, ttl)
             if auth_scheme == 'Bearer':
                 jwt_token = auth_header.split(' ')[1]
@@ -510,7 +584,7 @@ def _get_scheme_and_token(self, headers, http_method, uri, body, scopes,
         return None, None
 
 
-class OAuthClientCredentialsRequestAuthenticator(RequestAuthenticator):
+class OAuthClientCredentialsRequestAuthenticator(JWTRequestAuthenticator):
     """This class should authenticate OAuth2 client credentials grant
     type requests only. It will handle authenticating a request based
     on API key credentials.
@@ -523,7 +597,7 @@ def _get_scheme_and_token(self, headers, http_method, uri, body, scopes,
 
         if self._get_auth_scheme_from_header(auth_header) == auth_scheme:
             if body.get('grant_type') or url_qs.get('grant_type'):
-                return auth_scheme, _get_bearer_token(
+                return auth_scheme, self.jwt_factory.get_bearer_token(
                     self.app, scopes, http_method, uri, body, headers, ttl)
 
         return None, None

tests/mocks/test_api_auth.py

@@ -106,6 +106,70 @@ def test_basic_api_auth_with_generating_bearer_token(self):
         self.assertIsNotNone(result.token)
         self.assertEquals(result.token.scopes, ['test1'])
 
+    def test_basic_api_auth_with_generating_bearer_token_extra_data(self):
+        app = MagicMock()
+        app._client.auth.secret = 'fakeApiKeyProperties.secret'
+        app.href = 'HREF'
+        api_keys = MagicMock()
+        api_keys.get_key = lambda k, s=None: MagicMock(
+                id=FAKE_CLIENT_ID, secret=FAKE_CLIENT_SECRET, status=StatusMixin.STATUS_ENABLED)
+        app.api_keys = api_keys
+
+        basic_auth = base64.b64encode("{}:{}".format(FAKE_CLIENT_ID, FAKE_CLIENT_SECRET).encode('utf-8'))
+
+        uri = 'https://example.com/get'
+        http_method = 'GET'
+        body = {'grant_type': 'client_credentials', 'scope': 'test1'}
+        headers = {
+                'Authorization': b'Basic ' + basic_auth
+                }
+
+        allowed_scopes = ['test1']
+
+        authenticator = ApiRequestAuthenticator(app, token_data={"foo":"bar", "exp":"dumb"})
+        result = authenticator.authenticate(
+            headers=headers, http_method=http_method, uri=uri, body=body,
+            scopes=allowed_scopes)
+        self.assertIsNotNone(result)
+        self.assertIsNotNone(result.api_key)
+        self.assertIsNotNone(result.token)
+        self.assertEquals(result.token.claims['foo'], 'bar')
+        self.assertNotEquals(result.token.claims['exp'], 'dumb')
+        self.assertEquals(result.token.scopes, ['test1'])
+        # check that xsrfToken was not added
+        self.assertIsNone(result.token.claims.get("xsrfToken"))
+
+    def test_basic_api_auth_with_generating_bearer_token_xsrf(self):
+        app = MagicMock()
+        app._client.auth.secret = 'fakeApiKeyProperties.secret'
+        app.href = 'HREF'
+        api_keys = MagicMock()
+        api_keys.get_key = lambda k, s=None: MagicMock(
+                id=FAKE_CLIENT_ID, secret=FAKE_CLIENT_SECRET, status=StatusMixin.STATUS_ENABLED)
+        app.api_keys = api_keys
+
+        basic_auth = base64.b64encode("{}:{}".format(FAKE_CLIENT_ID, FAKE_CLIENT_SECRET).encode('utf-8'))
+
+        uri = 'https://example.com/get'
+        http_method = 'GET'
+        body = {'grant_type': 'client_credentials', 'scope': 'test1'}
+        headers = {
+                'Authorization': b'Basic ' + basic_auth
+                }
+
+        allowed_scopes = ['test1']
+
+        authenticator = ApiRequestAuthenticator(app, add_xsrf_token=True)
+        result = authenticator.authenticate(
+            headers=headers, http_method=http_method, uri=uri, body=body,
+            scopes=allowed_scopes)
+        self.assertIsNotNone(result)
+        self.assertIsNotNone(result.api_key)
+        self.assertIsNotNone(result.token)
+
+        # check that xsrfToken was added
+        uuid.UUID(result.token.claims.get("xsrfToken"))
+
     def test_basic_api_auth_with_invalid_scope_no_token_get_generated(self):
         app = MagicMock()
         app._client.auth.secret = 'fakeApiKeyProperties.secret'