Wendell 👾

Avatar of Wendell Hu

flask login 源码解析

July 12, 2017

There are currently no translations for this article.

Help us to translate it into a language that you are good with!

flask-login 的基本使用

在介绍 flask-login 的工作原理之前先来简要回顾一下 flask-login 的使用方法。

首先要创建一个 LoginManager 的实例并注册在 Flask 实例上,然后提供一个 user_loader 回调函数来根据会话中存储的用户 ID 来加载用户对象。

login_manager = LoginManager()
login_manager.init_app(app)

@login_manager.user_loader
def load_user(user_id):
    return User.get(user_id)

flask-login 还要求你对数据对象做一些改动,添加以下属性和方法:

@property
def is_active(self):
   return True

@property
def is_authenticated(self):
   return True

@property
def is_anonymous(self):
   return False

#: 这个方法返回一个能够识别唯一用户的 ID
def get_id(self):
   try:
       return text_type(self.id)
   except AttributeError:
       raise NotImplementedError('No `id` attribute - override `get_id`')

完成这些设置工作之后,就可以使用 flask-login 了,一些典型的用法包括登录和登出用户:login_user(user)logout_user(),及使用 @login_required 保护一些视图函数,检测当前用户是否有访问的权限(根据是否认证进行区别):

@app.route("/settings")
@login_required
def settings():
    pass

以及通过 current_user 对象来访问当前用户。

flask-login 源码解析

我们按照使用过程中调用 flask-login 的顺序来解析其源码。

LoginManager 对象

先来看 LoginManager 对象,它用于记录所有的配置信息,其 __init__ 方法中初始化了这些配置信息。一个 LoginManager 对象通过 init_app 方法注册到 Flask 实例上:

def init_app(self, app, add_context_processor=True):
   app.login_manager = self
   app.after_request(self._update_remember_cookie)

   self._login_disabled = app.config.get('LOGIN_DISABLED', False)

   if add_context_processor:
       app.context_processor(_user_context_processor)

这个方法的主要工作是在 Flask 实例的 afterrequest 钩子上添加了一个用户更新 rememberme cookie 的函数,并在 Flask 的上下文处理器中添加了一个用户上下文处理器。

def _user_context_processor():
    return dict(current_user=_get_user())

这个上下文处理器设置了一个全局可访问的变量 current_user,这样我们就可以在视图函数或者模板文件中访问这个变量了。

user_loader 修饰器

然后就到了这个方法,它是 LoginManager 的实例方法,把 usercallback 设置成我们传入的函数,在实际的使用过程中,我们是通过修饰器传入这个函数的,就是 `loaduser(user_id)` 函数。

def user_loader(self, callback):
   self.user_callback = callback
   return callback

该方法要求你的回调函数必须能够接收一个 unicode 编码的 ID 并返回一个用户对象,如果用户不存在就返回 None。

login_user 方法

我们跳过对 User 类的修改,直接来看这个方法。

def login_user(user, remember=False, force=False, fresh=True):
    if not force and not user.is_active:
        return False

    user_id = getattr(user, current_app.login_manager.id_attribute)()
    session['user_id'] = user_id
    session['_fresh'] = fresh
    session['_id'] = current_app.login_manager._session_identifier_generator()

    if remember:
        session['remember'] = 'set'

    _request_ctx_stack.top.user = user
    user_logged_in.send(current_app._get_current_object(), user=_get_user())
    return True

如果用户不活跃 not.is_active 而且不要求强制登录 force,就返回失败。否则,先得到 user_id,它是通过 getattr 函数访问 userlogin_manager.id_attribute 属性得到的。追根溯源,最终 getattr 访问的是 userget_id 方法,这就是为什么 flask-login 要求我们在 User 类中添加该方法。

然后在 Flask 提供的 session 中添加以下三个 session:user_id _fresh _id,其中 _id 是通过 LoginManager_session_identifier_generator 方法获取到的,而这个方法默认绑定在这个方法上:

def _create_identifier():
    user_agent = request.headers.get('User-Agent')
    if user_agent is not None:
        user_agent = user_agent.encode('utf-8')
    base = '{0}|{1}'.format(_get_remote_addr(), user_agent)
    if str is bytes:
        base = text_type(base, 'utf-8', errors='replace')  # pragma: no cover
    h = sha512()
    h.update(base.encode('utf8'))
    return h.hexdigest()

不用太深究,知道这个方法最终根据放着用户代理和 IP 信息生成了一个加盐的 ID 就行了,它的作用是防止有人伪造 cookie。

然后根据是否需要记住用户添加 remember session。最后,在 _request_ctx_stack.top 中添加该用户,发出一个用户登录信号后返回成功。在这个登录信号中,调用了 _get_user 方法,_get_user 方法的细节是先检测在 _request_ctx_stack.top 中有没有用户信息,如果没有,就通过 _load_user 方法在栈顶添加用户信息,如果有就返回这个用户对象。_load_user 方法很重要,但是在这里不会被调用,很明显 _request_ctx_stack.top 中肯定有 user 值,我们待会再来看这个方法。

def _get_user():
    if has_request_context() and not hasattr(_request_ctx_stack.top, 'user'):
        current_app.login_manager._load_user()

    return getattr(_request_ctx_stack.top, 'user', None)

login_required 修饰器

这个修饰器常被用来保护只有登录用户才能访问的视图函数,它会在实际调用视图函数之前先检查当前用户是否已经登录并认证,如果没有,就调用 LoginManager.unauthorized 这个回调函数,它还对一些 HTTP 方法和测试情况提供了例外处理。

def login_required(func):
    @wraps(func)
    def decorated_view(*args, **kwargs):
        if request.method in EXEMPT_METHODS:
            return func(*args, **kwargs)
        elif current_app.login_manager._login_disabled:
            return func(*args, **kwargs)
        elif not current_user.is_authenticated:
            return current_app.login_manager.unauthorized()
        return func(*args, **kwargs)
    return decorated_view

current_user 对象

在之前的分析中,可以看到这个变量经常出现并大有用途,开发者可以通过访问这个变量来获取到当前用户,如果用户未登录,获取到的就是一个匿名用户,它的定义:

current_user = LocalProxy(lambda: _get_user())

_get_user() 方法之前已经讲过,我们直接跳到 _load_user 方法。显然,如果用户登录后再次发出了请求,我们就要从 cookie,或者说,Flask 在此之上封装的 session 中获取用户信息才能正确地进行后续处理,_load_user 方法的作用就是这个,该方法如下:

def _load_user(self):
   user_accessed.send(current_app._get_current_object())

   config = current_app.config
   if config.get('SESSION_PROTECTION', self.session_protection):
       deleted = self._session_protection()
       if deleted:
           return self.reload_user()

   is_missing_user_id = 'user_id' not in session
   if is_missing_user_id:
       cookie_name = config.get('REMEMBER_COOKIE_NAME', COOKIE_NAME)
       header_name = config.get('AUTH_HEADER_NAME', AUTH_HEADER_NAME)
       has_cookie = (cookie_name in request.cookies and
                     session.get('remember') != 'clear')
       if has_cookie:
           return self._load_from_cookie(request.cookies[cookie_name])
       elif self.request_callback:
           return self._load_from_request(request)
       elif header_name in request.headers:
           return self._load_from_header(request.headers[header_name])

   return self.reload_user()

def _session_protection(self):
   sess = session._get_current_object()
   ident = self._session_identifier_generator()

   app = current_app._get_current_object()
   mode = app.config.get('SESSION_PROTECTION', self.session_protection)

   if sess and ident != sess.get('_id', None):
       if mode == 'basic' or sess.permanent:
           sess['_fresh'] = False
           session_protected.send(app)
           return False
       elif mode == 'strong':
           for k in SESSION_KEYS:
               sess.pop(k, None)

           sess['remember'] = 'clear'
           session_protected.send(app)
           return True

   return False

该方法首先保证 session 的安全,如果 session 通过了安全验证,就通过 reload_user 方法重载用户,否则检查 session 中是否没有 user_id 来重载用户,如果没有,通过三种不同的方式重载用户。

def reload_user(self, user=None):
    ctx = _request_ctx_stack.top

    if user is None:
        user_id = session.get('user_id')
        if user_id is None:
            ctx.user = self.anonymous_user()
        else:
            if self.user_callback is None:
                raise Exception(
                    "No user_loader has been installed for this "
                    "LoginManager. Add one with the "
                    "'LoginManager.user_loader' decorator.")
            user = self.user_callback(user_id)
            if user is None:
                ctx.user = self.anonymous_user()
            else:
                ctx.user = user
    else:
        ctx.user = user

在这个重载方法中,如果 user_id 不存在,就把匿名用户加载到 _request_ctx_stack.top,否则根据 user_id 加载用户,若该用户不存在,仍加载匿用户。

之后,current_user 就能获取到用户对象,或者是一个匿名用户对象了。

current_user = LocalProxy(lambda: _get_user())

logout_user 方法

这个方法先获取当前用户,然后移除 user_id _fresh 等 session,然后移除 remember,最后重载当前用户,很明显,重载之后会是一个匿名用户。

def logout_user():
    user = _get_user()

    if 'user_id' in session:
        session.pop('user_id')

    if '_fresh' in session:
        session.pop('_fresh')

    cookie_name = current_app.config.get('REMEMBER_COOKIE_NAME', COOKIE_NAME)
    if cookie_name in request.cookies:
        session['remember'] = 'clear'

    user_logged_out.send(current_app._get_current_object(), user=user)

    current_app.login_manager.reload_user()
    return True

remember_me cookie

记得我们之前提到 flask-login 在 Flask 实例的 afterrequest 钩子上添加了一个用户更新 rememberme cookie 的函数吗,我们显然需要在请求的最后对 remember 进行处理。

def _update_remember_cookie(self, response):
   # Don't modify the session unless there's something to do.
   if 'remember' in session:
       operation = session.pop('remember', None)
       if operation == 'set' and 'user_id' in session:
           self._set_cookie(response)
       elif operation == 'clear':
           self._clear_cookie(response)
    return response

这个函数根据是否要设置 remember 来调用不同的函数

def _set_cookie(self, response):
    config = current_app.config
    cookie_name = config.get('REMEMBER_COOKIE_NAME', COOKIE_NAME)
    duration = config.get('REMEMBER_COOKIE_DURATION', COOKIE_DURATION)
    domain = config.get('REMEMBER_COOKIE_DOMAIN')
    path = config.get('REMEMBER_COOKIE_PATH', '/')

    secure = config.get('REMEMBER_COOKIE_SECURE', COOKIE_SECURE)
    httponly = config.get('REMEMBER_COOKIE_HTTPONLY', COOKIE_HTTPONLY)

    data = encode_cookie(text_type(session['user_id']))

    try:
        expires = datetime.utcnow() + duration
    except TypeError:
        raise Exception('REMEMBER_COOKIE_DURATION must be a ' +
                        'datetime.timedelta, instead got: {0}'.format(
                            duration))

    response.set_cookie(cookie_name,
                        value=data,
                        expires=expires,
                        domain=domain,
                        path=path,
                        secure=secure,
                        httponly=httponly)

def _clear_cookie(self, response):
    config = current_app.config
    cookie_name = config.get('REMEMBER_COOKIE_NAME', COOKIE_NAME)
    domain = config.get('REMEMBER_COOKIE_DOMAIN')
    path = config.get('REMEMBER_COOKIE_PATH', '/')
    response.delete_cookie(cookie_name, domain=domain, path=path)

总结

  • flask-login 使用 Flask 提供的 session 来保存用户信息,通过 user_id 来记录用户身份,_id 来防止攻击者对 session 的伪造。
  • 通过 _request_ctx_stack.top.user,flask-login 实现了线程安全。
  • 通过 cookie 来实现 remember 功能。

其他功能如 fresh login 请自行查看源码了解。

仿造 flask-login 写一个基于 token 的身份认证模块

flask-login 虽然好用,但由于其是基于 session 的,对于无状态的 RESTful API 应用无能为力。我在一个最近的项目模仿了它的接口,实现了一个简单但是好用的身份认证模块。

from functools import wraps
from flask import (_request_ctx_stack, has_request_context, request,
                   current_app)
from flask_restful import abort
from werkzeug.local import LocalProxy
from app.models.user import User

#: a proxy for the current user
#: it would be an anonymous user if no user is logged in
current_user = LocalProxy(lambda: _get_user())


class AnonymousUserMixin(object):
    @property
    def is_active(self):
        return False

    @property
    def is_authenticated(self):
        return False

    @property
    def is_anonymous(self):
        return True

    def __repr__(self):
        return ''


class Manager(object):
    def __init__(self, app=None):
        if app:
            self.init_app(app)

    def init_app(self, app):
        app.login_manager = self
        app.context_processor(_user_context_processor)

        self._anonymous_user = AnonymousUserMixin
        self._login_disabled = app.config['LOGIN_DISABLED'] or False

    @staticmethod
    def _load_user():
        """Try to load user from request.json.token and set it to
        `_request_ctx_stack.top.user`. If None, set current user as an anonymous
        user.
        """
        ctx = _request_ctx_stack.top
        json = request.json
        user = AnonymousUserMixin()

        if json and json.get('token'):
            real_user = User.load_user_from_auth_token(json.get('token'))
            if real_user:
                user = real_user

        ctx.user = user


def _get_user():
    """Get current user from request context."""
    if has_request_context() and not hasattr(_request_ctx_stack.top, 'user'):
        current_app.login_manager._load_user()

    return getattr(_request_ctx_stack.top, 'user', None)


def _user_context_processor():
    """A context processor to prepare current user."""
    return dict(current_user=_get_user())


def login_user(user):
    """Login a user and return a token."""
    _request_ctx_stack.top.user = user
    return user.generate_auth_token()


def logout_user(user):
    """For a restful API there shouldn't be a `logout` method because the
    server is stateless.
    """
    pass


def login_required(func):
    """Decorator to protect view functions that should only be accessed
    by authenticated users.
    """

    @wraps(func)
    def decorated_view(*args, **kwargs):
        if current_app.login_manager._login_disabled:
            return func(*args, **kwargs)
        elif not current_user.is_authenticated:
            abort(403, err='40300',
                  message='Please login before carrying out this action.')
        return func(*args, **kwargs)

    return decorated_view

Wendell Hu

Written by Wendell Hu
Follow me on Twitter