Source

garter / garter / csrf.py

Full commit
from uuid import uuid4
from flask import abort, request, session, g
from werkzeug.routing import NotFound


_exempt_views = []

def csrf_exempt(view):
    _exempt_views.append(view)
    return view

def csrf(app, on_csrf=None):
    @app.before_request
    def _csrf_check_exemptions():
        try:
            dest = app.view_functions[app.match_request()[0]]
            g._csrf_exempt = dest in _exempt_views
        except NotFound:
            g._csrf_exempt = False
    
    @app.before_request
    def _csrf_protect():
        if not g._csrf_exempt:
            if request.method == "POST":
                csrf_token = session.pop('_csrf_token', None)
                if not csrf_token or csrf_token != request.form.get('_csrf_token'):
                    if on_csrf:
                        on_csrf(*app.match_request())
                    abort(400)
    
    def generate_csrf_token():
        if '_csrf_token' not in session:
            session['_csrf_token'] = str(uuid4())
        return session['_csrf_token']
    
    app.jinja_env.globals['csrf_token'] = generate_csrf_token