From 27ffe6dc3cfefc488749d8bd631ec229e1e42550 Mon Sep 17 00:00:00 2001
From: Alexander Pantyukhin <apantykhin@gmail.com>
Date: Thu, 26 Apr 2018 00:52:36 +0400
Subject: [PATCH] init for jwt identity (#147)

---
 aiohttp_security/jwt_identity.py | 34 +++++++++++++++++++
 requirements-dev.txt             |  2 ++
 tests/test_jwt_identity.py       | 56 ++++++++++++++++++++++++++++++++
 3 files changed, 92 insertions(+)
 create mode 100644 aiohttp_security/jwt_identity.py
 create mode 100644 tests/test_jwt_identity.py

diff --git a/aiohttp_security/jwt_identity.py b/aiohttp_security/jwt_identity.py
new file mode 100644
index 0000000..f51105e
--- /dev/null
+++ b/aiohttp_security/jwt_identity.py
@@ -0,0 +1,34 @@
+"""Identity policy for storing info in the jwt token.
+
+"""
+
+from .abc import AbstractIdentityPolicy
+try:
+    import jwt
+except ImportError:  # pragma: no cover
+    jwt = None
+
+
+AUTH_HEADER_NAME = 'Authorization'
+
+
+class JWTIdentityPolicy(AbstractIdentityPolicy):
+    def __init__(self, secret, algorithm=None):
+        if jwt is None:
+            raise RuntimeError("Please install pyjwt")
+        self.secret = secret
+        self.algorithm = 'HS256' if algorithm is None else algorithm
+
+    async def identify(self, request):
+        header_identity = request.headers.get(AUTH_HEADER_NAME)
+        identity = jwt.decode(header_identity,
+                              self.secret,
+                              algorithm=self.algorithm)
+
+        return identity['identity']
+
+    async def remember(self, *args, **kwargs):  # pragma: no cover
+        pass
+
+    async def forget(self, request, response):  # pragma: no cover
+        pass
diff --git a/requirements-dev.txt b/requirements-dev.txt
index 4fa4129..ac3dfa4 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -2,6 +2,7 @@
 flake8==3.5.0
 pytest==3.5.0
 pytest-cov==2.5.1
+pytest-mock==1.6.3
 coverage==4.5.1
 sphinx==1.7.3
 pep257==0.7.0
@@ -13,3 +14,4 @@ passlib==1.7.1
 cryptography==2.2.2
 aiohttp==3.1.3
 pytest-aiohttp==0.3.0
+pyjwt
\ No newline at end of file
diff --git a/tests/test_jwt_identity.py b/tests/test_jwt_identity.py
new file mode 100644
index 0000000..7373aa9
--- /dev/null
+++ b/tests/test_jwt_identity.py
@@ -0,0 +1,56 @@
+import pytest
+from aiohttp import web
+from aiohttp_security import AbstractAuthorizationPolicy
+from aiohttp_security import setup as _setup
+from aiohttp_security.jwt_identity import JWTIdentityPolicy
+from aiohttp_security.api import IDENTITY_KEY
+import jwt
+
+
+class Autz(AbstractAuthorizationPolicy):
+
+    async def permits(self, identity, permission, context=None):
+        pass
+
+    async def authorized_userid(self, identity):
+        pass
+
+
+async def test_no_pyjwt_installed(mocker):
+    mocker.patch('aiohttp_security.jwt_identity.jwt', None)
+    with pytest.raises(RuntimeError):
+        JWTIdentityPolicy('secret')
+
+
+async def test_identify(loop, test_client):
+    kwt_secret_key = 'Key'
+
+    async def create(request):
+        response = web.Response()
+        data = await request.post()
+
+        encoded_identity = jwt.encode({'identity': data['login']},
+                                      kwt_secret_key,
+                                      algorithm='HS256')
+
+        response.text = encoded_identity.decode('utf-8')
+        return response
+
+    async def check(request):
+        policy = request.app[IDENTITY_KEY]
+        user_id = await policy.identify(request)
+        assert 'Andrew' == user_id
+        return web.Response()
+
+    app = web.Application(loop=loop)
+    _setup(app, JWTIdentityPolicy(kwt_secret_key), Autz())
+    app.router.add_route('GET', '/', check)
+    app.router.add_route('POST', '/', create)
+    client = await test_client(app)
+    resp = await client.post('/', data={'login': 'Andrew'})
+    jwt_token = await resp.content.read()
+    assert 200 == resp.status
+    await resp.release()
+    headers = {'Authorization': str(jwt_token.decode('utf-8'))}
+    resp = await client.get('/', headers=headers)
+    assert 200 == resp.status