User sessions

miguelgrinberg · Aug 6, 2022 · 355ffef · 355ffef
1 parent 199d23f
commit 355ffef
Show file tree

Hide file tree

Showing 5 changed files with 186 additions and 0 deletions.
diff --git a/examples/login.py b/examples/login.py
@@ -0,0 +1,58 @@
+from microdot import Microdot, Response, redirect
+from microdot_session import set_session_secret_key, with_session, \
+    update_session, delete_session
+
+BASE_TEMPLATE = '''<!doctype html>
+<html>
+  <head>
+    <title>Microdot login example</title>
+  </head>
+  <body>
+    <h1>Microdot login example</h1>
+    {content}
+  </body>
+</html>'''
+
+LOGGED_OUT = '''<p>You are not logged in.</p>
+<form method="POST">
+  <p>
+    Username:
+    <input type="text" name="username" autofocus />
+  </p>
+  <input type="submit" value="Submit" />
+</form>'''
+
+LOGGED_IN = '''<p>Hello <b>{username}</b>!</p>
+<form method="POST" action="/logout">
+  <input type="submit" value="Logout" />
+</form>'''
+
+app = Microdot()
+set_session_secret_key('top-secret')
+Response.default_content_type = 'text/html'
+
+
+@app.get('/')
+@app.post('/')
+@with_session
+def index(req, session):
+    username = session.get('username')
+    if req.method == 'POST':
+        username = req.form.get('username')
+        update_session(req, {'username': username})
+        return redirect('/')
+    if username is None:
+        return BASE_TEMPLATE.format(content=LOGGED_OUT)
+    else:
+        return BASE_TEMPLATE.format(content=LOGGED_IN.format(
+            username=username))
+
+
+@app.post('/logout')
+def logout(req):
+    delete_session(req)
+    return redirect('/')
+
+
+if __name__ == '__main__':
+    app.run()
diff --git a/libs/common/ujwt.py → libs/micropython/jwt.py b/libs/common/ujwt.py → libs/micropython/jwt.py
diff --git a/src/microdot_session.py b/src/microdot_session.py
@@ -0,0 +1,55 @@
+import jwt
+
+secret_key = None
+
+
+def set_session_secret_key(key):
+    global secret_key
+    secret_key = key
+
+
+def get_session(request):
+    global secret_key
+    if not secret_key:
+        raise ValueError('The session secret key is not configured')
+    session = request.cookies.get('session')
+    if session is None:
+        return {}
+    try:
+        session = jwt.decode(session, secret_key, algorithms=['HS256'])
+    except jwt.exceptions.PyJWTError:  # pragma: no cover
+        raise
+        return {}
+    return session
+
+
+def update_session(request, session):
+    if not secret_key:
+        raise ValueError('The session secret key is not configured')
+
+    encoded_session = jwt.encode(session, secret_key, algorithm='HS256')
+
+    @request.after_request
+    def _update_session(request, response):
+        response.set_cookie('session', encoded_session, http_only=True)
+        return response
+
+
+def delete_session(request):
+    @request.after_request
+    def _delete_session(request, response):
+        response.set_cookie('session', '', http_only=True,
+                            expires='Thu, 01 Jan 1970 00:00:01 GMT')
+        return response
+
+
+def with_session(f):
+    def wrapper(request, *args, **kwargs):
+        return f(request, get_session(request), *args, **kwargs)
+
+    for attr in ['__name__', '__doc__', '__module__', '__qualname__']:
+        try:
+            setattr(wrapper, attr, getattr(f, attr))
+        except AttributeError:  # pragma: no cover
+            pass
+    return wrapper
diff --git a/tests/test_session.py b/tests/test_session.py
@@ -0,0 +1,72 @@
+import unittest
+from microdot import Microdot
+from microdot_session import set_session_secret_key, get_session, \
+    update_session, delete_session, with_session
+from microdot_test_client import TestClient
+
+set_session_secret_key('top-secret!')
+
+
+class TestSession(unittest.TestCase):
+    def setUp(self):
+        self.app = Microdot()
+        self.client = TestClient(self.app)
+
+    def tearDown(self):
+        pass
+
+    def test_session(self):
+        @self.app.get('/')
+        def index(req):
+            session = get_session(req)
+            return str(session.get('name'))
+
+        @self.app.get('/with')
+        @with_session
+        def session_context_manager(req, session):
+            return str(session.get('name'))
+
+        @self.app.post('/set')
+        def set_session(req):
+            update_session(req, {'name': 'joe'})
+            return 'OK'
+
+        @self.app.post('/del')
+        def del_session(req):
+            delete_session(req)
+            return 'OK'
+
+        res = self.client.get('/')
+        self.assertEqual(res.text, 'None')
+        res = self.client.get('/with')
+        self.assertEqual(res.text, 'None')
+
+        res = self.client.post('/set')
+        self.assertEqual(res.text, 'OK')
+
+        res = self.client.get('/')
+        self.assertEqual(res.text, 'joe')
+        res = self.client.get('/with')
+        self.assertEqual(res.text, 'joe')
+
+        res = self.client.post('/del')
+        self.assertEqual(res.text, 'OK')
+
+        res = self.client.get('/')
+        self.assertEqual(res.text, 'None')
+        res = self.client.get('/with')
+        self.assertEqual(res.text, 'None')
+
+    def test_session_no_secret_key(self):
+        set_session_secret_key(None)
+
+        @self.app.get('/')
+        def index(req):
+            self.assertRaises(ValueError, get_session, req)
+            self.assertRaises(ValueError, update_session, req, {})
+            return ''
+
+        res = self.client.get('/')
+        self.assertEqual(res.status_code, 200)
+
+        set_session_secret_key('top-secret!')
diff --git a/tox.ini b/tox.ini
@@ -20,6 +20,7 @@ deps=
     pytest
     pytest-cov
     jinja2
+    pyjwt
 setenv=
     PYTHONPATH=libs/common