1 files changed, 17 insertions, 2 deletions
diff --git a/auth.py b/auth.py
index 5a4790d..7f330db 100644
--- a/auth.py
+++ b/auth.py
@@ -1,11 +1,13 @@
 from flask_peewee.auth import Auth
-from flask import session
+from flask_peewee.utils import get_next
+from flask import session, url_for, request, redirect
 from models import User, Session
 from app import app, db, pad
 from datetime import datetime
 from padlite import APIException
 import ldap
 import uuid
+import functools
 
 class LdapAuth(Auth):
     def get_user_model(self):
@@ -70,4 +72,17 @@ class LdapAuth(Auth):
 
         return reduce(escape, chars_to_escape, s)
 
-auth = LdapAuth(app, db, user_model=User)
+    def test_user(self, test_fn):
+        def decorator(fn):
+            @functools.wraps(fn)
+            def inner(*args, **kwargs):
+                user = self.get_logged_in_user()
+
+                if not user or not test_fn(user):
+                    login_url = url_for('%s.login' % self.blueprint.name, next="%s%s" % (request.environ['SCRIPT_NAME'], get_next()))
+                    return redirect(login_url)
+                return fn(*args, **kwargs)
+            return inner
+        return decorator
+
+auth = LdapAuth(app, db, user_model=User, default_next_url='/teams')