Shotca avatar Shotca committed 2fdc837

First steps in reworking to not pickle the user but the user.__dict__

Comments (0)

Files changed (9)

examples/sqlalchemy_model.py

-import datetime
 from flask import Flask, request, redirect, url_for
 from flaskext.sqlalchemy import SQLAlchemy
-from flaskext.auth import Auth, AuthUser, login_required, logout
+from flaskext.auth import Auth, login_required, logout
+from flaskext.auth.models.sa import get_user_class
 
 app = Flask(__name__)
 app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:////tmp/test.db'
 db = SQLAlchemy(app)
 auth = Auth(app, login_url_name='index')
 
-def now():
-    return datetime.datetime.utcnow()
-
-class User(db.Model, AuthUser):
-    """
-    Implementation of User for SQLAlchemy.
-    """
-    id = db.Column(db.Integer, primary_key=True)
-    username = db.Column(db.String(80), unique=True, nullable=False)
-    password = db.Column(db.String(120), nullable=False)
-    salt = db.Column(db.String(80))
-    role = db.Column(db.String(80))
-    created = db.Column(db.DateTime(), default=now)
-    modified = db.Column(db.DateTime())
-
-    def __init__(self, *args, **kwargs):
-        super(User, self).__init__(*args, **kwargs)
-        password = kwargs.get('password')
-        if password is not None and not self.id:
-            self.created = datetime.datetime.utcnow()
-            # Initialize and encrypt password before first save.
-            self.set_and_encrypt_password(password)
+User = get_user_class(db.Model)
 
 @login_required()
 def admin():

flaskext/auth/__init__.py

     :license: BSD, see LICENSE for more details.
 """
 
-from flaskext.auth.auth import Auth, AuthUser, login, logout, get_current_user,\
-        login_required, encrypt
+from flaskext.auth.auth import Auth, AuthUser, login, logout, \
+        get_current_user_data, login_required, encrypt
 from flaskext.auth.permissions import has_permission, permission_required, \
         Role, Permission

flaskext/auth/auth.py

 
     role = None
 
-    def __init__(self, username, password=None, salt=None, role=None):
+    def __init__(self, username=None, password=None, salt=None, role=None):
         self.username = username
         # Storing password unmodified. Encryption of the password should 
         # happen explicitly.
             return True
         return False
 
-    def __eq__(self, other):
-        return hasattr(other, 'username') and self.username == other.username
+    @classmethod
+    def load_from_user_data(cls, user_data):
+        user = cls()
+        user.__dict__ = user_data
+        return user
 
     def is_logged_in(self):
-        return get_current_user() == self
+        user_data = get_current_user_data()
+        return user_data is not None and user_data.get('username') == self.username
 
 def encrypt(password, salt=None, hash_algorithm=None):
     """Encrypts a password based on the hashing algorithm."""
     Logs the user in. Note that NO AUTHENTICATION is done by this function. If
     you want to authenticate a user, use the AuthUser.authenticate() method.
     """
-    session[SESSION_USER_KEY] = user
+    session[SESSION_USER_KEY] = user.__dict__
     session[SESSION_LOGIN_KEY] = datetime.datetime.utcnow()
 
 def logout():
-    """Logs the currently logged in user out and returns the user (if any)."""
+    """Logs the currently logged in user out and returns the user data."""
     session.pop(SESSION_LOGIN_KEY, None)
     return session.pop(SESSION_USER_KEY, None)
     
-def get_current_user(apply_timeout=True):
+def get_current_user_data(apply_timeout=True):
+    """ 
+    Returns the data of the current user (user.__dict__) if there is a
+    current user and he didn't time out yet. If timeout should be ignored,
+    provide apply_timeout=False.  
     """
-    Returns the current user if there is one and he didn't time out yet. If
-    timeout should be ignored, provide apply_timeout=False.
-    """
-    user = session.get(SESSION_USER_KEY, None)
-    if user is None:
-        return None
+    user_data = session.get(SESSION_USER_KEY, None)
+    if user_data is None:
+        return None 
     if not apply_timeout:
-        return user
+        return user_data
     login_datetime = session[SESSION_LOGIN_KEY]
     now = datetime.datetime.utcnow()
     user_timeout = current_app.auth.user_timeout
        datetime.timedelta(seconds=user_timeout):
         logout()
         return None
-    return user
+    return user_data
 
 def _not_logged_in(callback, *args, **kwargs):
     if callback is None:
     """
     def wrap(func):
         def decorator(*args, **kwargs):
-            if get_current_user() is None:
+            if get_current_user_data() is None:
                 return _not_logged_in(callback, *args, **kwargs)
             return func(*args, **kwargs)
         return decorator
Add a comment to this file

flaskext/auth/models/__init__.py

Empty file added.

Add a comment to this file

flaskext/auth/models/__init__.pyc

Binary file added.

flaskext/auth/models/sa.py

+import datetime
+from sqlalchemy import Column, Integer, String, DateTime
+from flaskext.auth import AuthUser
+
+def get_user_class(declarative_base):
+    class User(declarative_base, AuthUser):
+        """
+        Implementation of User for SQLAlchemy.
+        """
+        id = Column(Integer, primary_key=True)
+        username = Column(String(80), unique=True, nullable=False)
+        password = Column(String(120), nullable=False)
+        salt = Column(String(80))
+        role = Column(String(80))
+        created = Column(DateTime(), default=datetime.datetime.utcnow)
+        modified = Column(DateTime())
+
+        def __init__(self, *args, **kwargs):
+            super(User, self).__init__(*args, **kwargs)
+            password = kwargs.get('password')
+            if password is not None and not self.id:
+                self.created = datetime.datetime.utcnow()
+                # Initialize and encrypt password before first save.
+                self.set_and_encrypt_password(password)
+
+        def __getstate__(self):
+            global User
+            User = self.__class__
+            return (self.__dict__, self.__class__.__bases__[0])
+
+        def __setstate__(self, state):
+            self.__class__ = get_user_class(state[1])
+            self.__dict__.update(state[0])
+    return User
Add a comment to this file

flaskext/auth/models/sa.pyc

Binary file added.

flaskext/auth/permissions.py

 """
 
 from flask import current_app
-from flaskext.auth.auth import get_current_user, _not_logged_in
+from flaskext.auth.auth import AuthUser, get_current_user_data, _not_logged_in
 user_permissions = dict()
 
 def has_permission(user, resource, action):
     """
     def wrap(func):
         def decorator(*args, **kwargs):
-            if get_current_user() is None:
+            if get_current_user_data() is None:
                 return _not_logged_in(callback, *args, **kwargs)
-            if not has_permission(get_current_user(), resource, action):
+            if not has_permission(AuthUser.load_from_user_data(
+                    get_current_user_data()), resource, action):
                 if callback is None:
                     return current_app.auth.not_permitted_callback(*args, **kwargs)
                 else:
 import unittest, hashlib
-from flask import Flask, session, url_for
-from flaskext.auth import Auth, AuthUser, login, logout, get_current_user, \
+from flask import Flask, session
+from flaskext.auth import Auth, AuthUser, login, logout, get_current_user_data,\
         encrypt, Role, Permission, has_permission, login_required, \
         permission_required
 from flaskext.auth.auth import SESSION_USER_KEY, SESSION_LOGIN_KEY
     def test_login(self):
         with self.app.test_request_context():
             login(self.user)
-            assert session[SESSION_USER_KEY] == self.user
+            assert session[SESSION_USER_KEY] == self.user.__dict__
             assert session.get(SESSION_LOGIN_KEY)
 
     def test_current_user(self):
         with self.app.test_request_context():
             login(self.user)
-            assert get_current_user() == self.user
+            assert self.user.is_logged_in()
 
     def test_logout(self):
         with self.app.test_request_context():
             login(self.user)
-            user = logout()
-            assert user == self.user
+            user_data = logout()
+            assert user_data == self.user.__dict__
             assert session.get(SESSION_USER_KEY) is None
             assert session.get(SESSION_LOGIN_KEY) is None
 
             self.app.auth.user_timeout = 0.01
             login(self.user)
             time.sleep(0.02)
-            assert get_current_user() is None
+            assert get_current_user_data() is None
 
     def test_user_expiration_override(self):
         import time
             self.app.auth.user_timeout = 0.01
             login(self.user)
             time.sleep(0.02)
-            assert get_current_user(apply_timeout=False) == self.user
+            assert get_current_user_data(apply_timeout=False).\
+                    get('username') == self.user.username
 
     def test_authenticate(self):
         with self.app.test_request_context():
             assert self.user.authenticate(self.PASSWORD) == True
             assert self.user.is_logged_in() == True
-            assert get_current_user() == self.user
 
     def test_authenticate_fail(self):
         with self.app.test_request_context():
Tip: Filter by directory path e.g. /media app.js to search for public/media/app.js.
Tip: Use camelCasing e.g. ProjME to search for ProjectModifiedEvent.java.
Tip: Filter by extension type e.g. /repo .js to search for all .js files in the /repo directory.
Tip: Separate your search with spaces e.g. /ssh pom.xml to search for src/ssh/pom.xml.
Tip: Use ↑ and ↓ arrow keys to navigate and return to view the file.
Tip: You can also navigate files with Ctrl+j (next) and Ctrl+k (previous) and view the file with Ctrl+o.
Tip: You can also navigate files with Alt+j (next) and Alt+k (previous) and view the file with Alt+o.