Repository: pallets-eco/flask-wtf (python)

Issue:
Make it easier to access a CSRF token in automated tests
If you want to run your automated tests with CSRF enabled (which is a good idea if it's enabled in production), there's no good built-in way to do so. Even the tests for this project [use regular expressions to parse the CSRF token out of the page](https://github.com/lepture/flask-wtf/blob/3c9dcf5cc/tests/test_csrf.py#L10-L18), which is brittle and confusing. It would be better to provide some way to access the CSRF token in the Flask test client itself.

I've written a GitHub Gist that walks though how I implemented this myself, but maybe Flask-WTF could change some internals to make it cleaner and easier? https://gist.github.com/singingwolfboy/2fca1de64950d5dfed72



Discussion:

Patch:
diff --git a/flask_wtf/_compat.py b/flask_wtf/_compat.py
--- a/flask_wtf/_compat.py
+++ b/flask_wtf/_compat.py
@@ -6,9 +6,11 @@
 if not PY2:
     text_type = str
     string_types = (str,)
+    from urllib.parse import urlparse
 else:
     text_type = unicode
     string_types = (str, unicode)
+    from urlparse import urlparse
 
 
 def to_bytes(text):
diff --git a/flask_wtf/csrf.py b/flask_wtf/csrf.py
--- a/flask_wtf/csrf.py
+++ b/flask_wtf/csrf.py
@@ -8,128 +8,94 @@
     :copyright: (c) 2013 by Hsiaoming Yang.
 """
 
-import os
-import hmac
 import hashlib
-import time
-from flask import Blueprint
-from flask import current_app, session, request, abort
+import os
+import warnings
+from functools import wraps
+
+from flask import Blueprint, current_app, request, session
+from itsdangerous import BadData, URLSafeTimedSerializer
+from werkzeug.exceptions import BadRequest
 from werkzeug.security import safe_str_cmp
-from ._compat import to_bytes, string_types
-try:
-    from urlparse import urlparse
-except ImportError:
-    # python 3
-    from urllib.parse import urlparse
 
+from ._compat import FlaskWTFDeprecationWarning, string_types, urlparse
 
 __all__ = ('generate_csrf', 'validate_csrf', 'CsrfProtect')
 
 
-def generate_csrf(secret_key=None, time_limit=None, token_key='csrf_token', url_safe=False):
-    """Generate csrf token code.
-
-    :param secret_key: A secret key for mixing in the token,
-                       default is Flask.secret_key.
-    :param time_limit: Token valid in the time limit,
-                       default is 3600s.
-    """
+def _get_secret_key(secret_key=None):
     if not secret_key:
-        secret_key = current_app.config.get(
-            'WTF_CSRF_SECRET_KEY', current_app.secret_key
-        )
+        secret_key = current_app.config.get('WTF_CSRF_SECRET_KEY', current_app.secret_key)
 
     if not secret_key:
-        raise Exception('Must provide secret_key to use csrf.')
+        raise Exception('Must provide secret_key to use CSRF.')
+
+    return secret_key
 
-    if time_limit is None:
-        time_limit = current_app.config.get('WTF_CSRF_TIME_LIMIT', 3600)
 
-    if token_key not in session:
-        session[token_key] = hashlib.sha1(os.urandom(64)).hexdigest()
-
-    if time_limit:
-        expires = int(time.time() + time_limit)
-        csrf_build = '%s%s' % (session[token_key], expires)
-    else:
-        expires = ''
-        csrf_build = session[token_key]
-
-    hmac_csrf = hmac.new(
-        to_bytes(secret_key),
-        to_bytes(csrf_build),
-        digestmod=hashlib.sha1
-    ).hexdigest()
-    delimiter = '--' if url_safe else '##'
-    return '%s%s%s' % (expires, delimiter, hmac_csrf)
-
-
-def validate_csrf(data, secret_key=None, time_limit=None, token_key='csrf_token', url_safe=False):
-    """Check if the given data is a valid csrf token.
-
-    :param data: The csrf token value to be checked.
-    :param secret_key: A secret key for mixing in the token,
-                       default is Flask.secret_key.
-    :param time_limit: Check if the csrf token is expired.
-                       default is True.
+def generate_csrf(secret_key=None, token_key='csrf_token'):
+    """Generate a CSRF token. The token is cached for a request, so multiple
+    calls to this function will generate the same token.
+
+    During testing, it might be useful to access the signed token in
+    ``request.csrf_token`` and the raw token in ``session['csrf_token']``.
+
+    :param secret_key: Used to securely sign the token. Default is
+        ``WTF_CSRF_SECRET_KEY`` or ``SECRET_KEY``.
+    :param token_key: key where token is stored in session for comparision.
     """
-    delimiter = '--' if url_safe else '##'
-    if not data or delimiter not in data:
-        return False
 
-    try:
-        expires, hmac_csrf = data.split(delimiter, 1)
-    except ValueError:
-        return False  # unpack error
+    if not getattr(request, token_key, None):
+        if token_key not in session:
+            session[token_key] = hashlib.sha1(os.urandom(64)).hexdigest()
 
-    if time_limit is None:
-        time_limit = current_app.config.get('WTF_CSRF_TIME_LIMIT', 3600)
+        s = URLSafeTimedSerializer(_get_secret_key(secret_key), salt='wtf-csrf-token')
+        setattr(request, token_key, s.dumps(session[token_key]))
 
-    if time_limit:
-        try:
-            expires = int(expires)
-        except ValueError:
-            return False
+    return getattr(request, token_key)
 
-        now = int(time.time())
-        if now > expires:
-            return False
 
-    if not secret_key:
-        secret_key = current_app.config.get(
-            'WTF_CSRF_SECRET_KEY', current_app.secret_key
-        )
+def validate_csrf(data, secret_key=None, time_limit=None, token_key='csrf_token'):
+    """Check if the given data is a valid CSRF token. This compares the given
+    signed token to the one stored in the session.
 
-    if token_key not in session:
+    :param data: The signed CSRF token to be checked.
+    :param secret_key: Used to securely sign the token. Default is
+        ``WTF_CSRF_SECRET_KEY`` or ``SECRET_KEY``.
+    :param time_limit: Number of seconds that the token is valid. Default is
+        ``WTF_CSRF_TIME_LIMIT`` or 3600 seconds (60 minutes).
+    :param token_key: key where token is stored in session for comparision.
+    """
+
+    if not data or token_key not in session:
         return False
 
-    csrf_build = '%s%s' % (session[token_key], expires)
-    hmac_compare = hmac.new(
-        to_bytes(secret_key),
-        to_bytes(csrf_build),
-        digestmod=hashlib.sha1
-    ).hexdigest()
+    s = URLSafeTimedSerializer(_get_secret_key(secret_key), salt='wtf-csrf-token')
 
-    return safe_str_cmp(hmac_compare, hmac_csrf)
+    if time_limit is None:
+        time_limit = current_app.config.get('WTF_CSRF_TIME_LIMIT', 3600)
 
+    try:
+        token = s.loads(data, max_age=time_limit)
+    except BadData:
+        return False
 
-class CsrfProtect(object):
-    """Enable csrf protect for Flask.
+    return safe_str_cmp(session[token_key], token)
 
-    Register it with::
 
-        app = Flask(__name__)
-        CsrfProtect(app)
-
-    And in the templates, add the token input::
+class CsrfProtect(object):
+    """Enable CSRF protection globally for a Flask app.
 
-        <input type="hidden" name="csrf_token" value="{{ csrf_token() }}"/>
+    ::
 
-    If you need to send the token via AJAX, and there is no form::
+        app = Flask(__name__)
+        csrf = CsrfProtect(app)
 
-        <meta name="csrf_token" content="{{ csrf_token() }}" />
+    Checks the ``csrf_token`` field sent with forms, or the ``X-CSRFToken``
+    header sent with JavaScript requests. Render the token in templates using
+    ``{{ csrf_token() }}``.
 
-    You can grab the csrf token with JavaScript, and send the token together.
+    See the :ref:`csrf` documentation.
     """
 
     def __init__(self, app=None):
@@ -140,24 +106,19 @@ def __init__(self, app=None):
             self.init_app(app)
 
     def init_app(self, app):
-        self._app = app
-        app.jinja_env.globals['csrf_token'] = generate_csrf
-        app.config.setdefault(
-            'WTF_CSRF_HEADERS', ['X-CSRFToken', 'X-CSRF-Token']
-        )
-        app.config.setdefault('WTF_CSRF_SSL_STRICT', True)
         app.config.setdefault('WTF_CSRF_ENABLED', True)
         app.config.setdefault('WTF_CSRF_CHECK_DEFAULT', True)
-        app.config.setdefault('WTF_CSRF_METHODS', ['POST', 'PUT', 'PATCH'])
+        app.config['WTF_CSRF_METHODS'] = set(app.config.get(
+            'WTF_CSRF_METHODS', ['POST', 'PUT', 'PATCH', 'DELETE']
+        ))
+        app.config.setdefault('WTF_CSRF_HEADERS', ['X-CSRFToken', 'X-CSRF-Token'])
+        app.config.setdefault('WTF_CSRF_SSL_STRICT', True)
 
-        # expose csrf_token as a helper in all templates
-        @app.context_processor
-        def csrf_token():
-            return dict(csrf_token=generate_csrf)
+        app.jinja_env.globals['csrf_token'] = generate_csrf
+        app.context_processor(lambda: {'csrf_token': generate_csrf})
 
         @app.before_request
         def _csrf_protect():
-            # many things come from django.middleware.csrf
             if not app.config['WTF_CSRF_ENABLED']:
                 return
 
@@ -171,15 +132,17 @@ def _csrf_protect():
                 return
 
             view = app.view_functions.get(request.endpoint)
+
             if not view:
                 return
 
-            if self._exempt_views or self._exempt_blueprints:
-                dest = '%s.%s' % (view.__module__, view.__name__)
-                if dest in self._exempt_views:
-                    return
-                if request.blueprint in self._exempt_blueprints:
-                    return
+            if request.blueprint in self._exempt_blueprints:
+                return
+
+            dest = '%s.%s' % (view.__module__, view.__name__)
+
+            if dest in self._exempt_views:
+                return
 
             self.protect()
 
@@ -190,85 +153,119 @@ def _get_csrf_token(self):
         for key in request.form:
             if key.endswith('csrf_token'):
                 csrf_token = request.form[key]
+
                 if csrf_token:
                     return csrf_token
 
-        for header_name in self._app.config['WTF_CSRF_HEADERS']:
+        for header_name in current_app.config['WTF_CSRF_HEADERS']:
             csrf_token = request.headers.get(header_name)
+
             if csrf_token:
                 return csrf_token
+
         return None
 
     def protect(self):
-        if request.method not in self._app.config['WTF_CSRF_METHODS']:
+        if request.method not in current_app.config['WTF_CSRF_METHODS']:
             return
 
         if not validate_csrf(self._get_csrf_token()):
-            reason = 'CSRF token missing or incorrect.'
-            return self._error_response(reason)
+            self._error_response('CSRF token missing or incorrect.')
 
-        if request.is_secure and self._app.config['WTF_CSRF_SSL_STRICT']:
+        if request.is_secure and current_app.config['WTF_CSRF_SSL_STRICT']:
             if not request.referrer:
-                reason = 'Referrer checking failed - no Referrer.'
-                return self._error_response(reason)
+                self._error_response('Referrer checking failed - no Referrer.')
 
             good_referrer = 'https://%s/' % request.host
+
             if not same_origin(request.referrer, good_referrer):
-                reason = 'Referrer checking failed - origin does not match.'
-                return self._error_response(reason)
+                self._error_response('Referrer checking failed - origin does not match.')
 
         request.csrf_valid = True  # mark this request is csrf valid
 
     def exempt(self, view):
-        """A decorator that can exclude a view from csrf protection.
-
-        Remember to put the decorator above the `route`::
+        """Mark a view or blueprint to be excluded from CSRF protection.
 
-            csrf = CsrfProtect(app)
+        ::
 
-            @csrf.exempt
             @app.route('/some-view', methods=['POST'])
+            @csrf.exempt
             def some_view():
-                return
+                ...
+
+        ::
+
+            bp = Blueprint(...)
+            csrf.exempt(bp)
+
         """
+
         if isinstance(view, Blueprint):
             self._exempt_blueprints.add(view.name)
             return view
+
         if isinstance(view, string_types):
             view_location = view
         else:
             view_location = '%s.%s' % (view.__module__, view.__name__)
+
         self._exempt_views.add(view_location)
         return view
 
     def _error_response(self, reason):
-        return abort(400, reason)
+        raise CsrfError(reason)
 
     def error_handler(self, view):
-        """A decorator that set the error response handler.
+        """Register a function that will generate the response for CSRF errors.
 
-        It accepts one parameter `reason`::
+        .. deprecated:: 0.14
+            Use the standard Flask error system with
+            ``@app.errorhandler(CsrfError)`` instead. This will be removed in
+            version 1.0.
+
+        The function will be passed one argument, ``reason``. By default it will
+        raise a :class:`~flask_wtf.csrf.CsrfError`. ::
 
             @csrf.error_handler
             def csrf_error(reason):
                 return render_template('error.html', reason=reason)
 
-        By default, it will return a 400 response.
+        Due to historical reasons, the function may either return a response
+        or raise an exception with :func:`flask.abort`.
         """
-        self._error_response = view
+
+        warnings.warn(FlaskWTFDeprecationWarning(
+            '"@csrf.error_handler" is deprecated. Use the standard Flask error '
+            'system with "@app.errorhandler(CsrfError)" instead. This will be'
+            'removed in 1.0.'
+        ), stacklevel=2)
+
+        @wraps(view)
+        def handler(reason):
+            response = current_app.make_response(view(reason))
+            raise CsrfError(response.get_data(as_text=True), response=response)
+
+        self._error_response = handler
         return view
 
 
-def same_origin(current_uri, compare_uri):
-    parsed_uri = urlparse(current_uri)
-    parsed_compare = urlparse(compare_uri)
+class CsrfError(BadRequest):
+    """Raise if the client sends invalid CSRF data with the request.
 
-    if parsed_uri.scheme != parsed_compare.scheme:
-        return False
+    Generates a 400 Bad Request response with the failure reason by default.
+    Customize the response by registering a handler with
+    :meth:`flask.Flask.errorhandler`.
+    """
 
-    if parsed_uri.hostname != parsed_compare.hostname:
-        return False
+    description = 'CSRF token missing or incorrect.'
 
-    if parsed_uri.port != parsed_compare.port:
-        return False
-    return True
+
+def same_origin(current_uri, compare_uri):
+    current = urlparse(current_uri)
+    compare = urlparse(compare_uri)
+
+    return (
+        current.scheme == compare.scheme
+        and current.hostname == compare.hostname
+        and current.port == compare.port
+    )
diff --git a/flask_wtf/form.py b/flask_wtf/form.py
--- a/flask_wtf/form.py
+++ b/flask_wtf/form.py
@@ -1,16 +1,14 @@
 # coding: utf-8
 import warnings
 
-import werkzeug.datastructures
-from flask import request, session, current_app
+from flask import current_app, request, session
 from jinja2 import Markup
-from wtforms.compat import with_metaclass
+from werkzeug.datastructures import MultiDict
 from wtforms.ext.csrf.form import SecureForm
-from wtforms.form import FormMeta
 from wtforms.validators import ValidationError
-from wtforms.widgets import HiddenInput, SubmitInput
+from wtforms.widgets import HiddenInput
 
-from ._compat import text_type, string_types, FlaskWTFDeprecationWarning
+from ._compat import FlaskWTFDeprecationWarning, string_types, text_type
 from .csrf import generate_csrf, validate_csrf
 
 try:
@@ -70,7 +68,7 @@ def __init__(self, formdata=_Auto, obj=None, prefix='', csrf_context=None,
                     formdata = formdata.copy()
                     formdata.update(request.files)
                 elif request.get_json():
-                    formdata = werkzeug.datastructures.MultiDict(request.get_json())
+                    formdata = MultiDict(request.get_json())
             else:
                 formdata = None
 
@@ -94,25 +92,24 @@ def __init__(self, formdata=_Auto, obj=None, prefix='', csrf_context=None,
     def generate_csrf_token(self, csrf_context=None):
         if not self.csrf_enabled:
             return None
-        return generate_csrf(self.SECRET_KEY, self.TIME_LIMIT)
+
+        return generate_csrf(secret_key=self.SECRET_KEY)
 
     def validate_csrf_token(self, field):
         if not self.csrf_enabled:
             return True
-        if hasattr(request, 'csrf_valid') and request.csrf_valid:
+
+        if getattr(request, 'csrf_valid', False):
             # this is validated by CsrfProtect
             return True
-        if not validate_csrf(field.data, self.SECRET_KEY, self.TIME_LIMIT):
+
+        if not self.validate_csrf_data(field.data):
             raise ValidationError(field.gettext('CSRF token missing'))
 
     def validate_csrf_data(self, data):
-        """Check if the csrf data is valid.
+        """Check if the given data is a valid CSRF token."""
 
-        .. versionadded: 0.9.0
-
-        :param data: the csrf string to be validated.
-        """
-        return validate_csrf(data, self.SECRET_KEY, self.TIME_LIMIT)
+        return validate_csrf(data, secret_key=self.SECRET_KEY, time_limit=self.TIME_LIMIT)
 
     def is_submitted(self):
         """Consider the form submitted if there is an active request and


Test patch:
diff --git a/tests/base.py b/tests/base.py
--- a/tests/base.py
+++ b/tests/base.py
@@ -1,10 +1,12 @@
 from __future__ import with_statement
 
-from flask import Flask, render_template, jsonify
-from wtforms import StringField, HiddenField, SubmitField
-from wtforms.validators import DataRequired
+from unittest import TestCase as _TestCase
+
+from flask import Flask, jsonify, render_template
 from flask_wtf import FlaskForm
 from flask_wtf._compat import text_type
+from wtforms import HiddenField, StringField, SubmitField
+from wtforms.validators import DataRequired
 
 
 def to_unicode(text):
@@ -37,7 +39,7 @@ class SimpleForm(FlaskForm):
     pass
 
 
-class TestCase(object):
+class TestCase(_TestCase):
     def setUp(self):
         self.app = self.create_app()
         self.client = self.app.test_client()
diff --git a/tests/templates/csrf_macro.html b/tests/templates/csrf_macro.html
--- a/tests/templates/csrf_macro.html
+++ b/tests/templates/csrf_macro.html
@@ -1,3 +1,3 @@
 {% macro render_csrf_token() %}
-  <input type="hidden" name="csrf_token" value="{{ csrf_token() }}">
+  <input name="csrf_token" type="hidden" value="{{ csrf_token() }}">
 {% endmacro %}
diff --git a/tests/test_csrf.py b/tests/test_csrf.py
--- a/tests/test_csrf.py
+++ b/tests/test_csrf.py
@@ -1,21 +1,13 @@
 from __future__ import with_statement
 
 import re
-from flask import Blueprint
-from flask import render_template
-from flask_wtf.csrf import CsrfProtect
-from flask_wtf.csrf import validate_csrf, generate_csrf
-from .base import TestCase, MyForm, to_unicode
+import warnings
 
-csrf_token_input = re.compile(
-    r'name="csrf_token" type="hidden" value="([0-9a-z#A-Z-\.]*)"'
-)
+from flask import Blueprint, abort, render_template, request
+from flask_wtf._compat import FlaskWTFDeprecationWarning
+from flask_wtf.csrf import CsrfError, CsrfProtect, generate_csrf, validate_csrf
 
-
-def get_csrf_token(data):
-    match = csrf_token_input.search(to_unicode(data))
-    assert match
-    return match.groups()[0]
+from .base import MyForm, TestCase
 
 
 class TestCSRF(TestCase):
@@ -59,9 +51,9 @@ def test_invalid_csrf(self):
         response = self.client.post("/", data={"name": "danny"})
         assert response.status_code == 400
 
-        @self.csrf.error_handler
-        def invalid(reason):
-            return reason
+        @self.app.errorhandler(CsrfError)
+        def handle_csrf_error(e):
+            return e, 200
 
         response = self.client.post("/", data={"name": "danny"})
         assert response.status_code == 200
@@ -86,8 +78,9 @@ def test_invalid_secure_csrf3(self):
         assert response.status_code == 400
 
     def test_valid_csrf(self):
-        response = self.client.get("/")
-        csrf_token = get_csrf_token(response.data)
+        with self.client:
+            self.client.get('/')
+            csrf_token = request.csrf_token
 
         response = self.client.post("/", data={
             "name": "danny",
@@ -96,8 +89,9 @@ def test_valid_csrf(self):
         assert b'DANNY' in response.data
 
     def test_prefixed_csrf(self):
-        response = self.client.get('/')
-        csrf_token = get_csrf_token(response.data)
+        with self.client:
+            self.client.get('/')
+            csrf_token = request.csrf_token
 
         response = self.client.post('/', data={
             'prefix-name': 'David',
@@ -106,8 +100,9 @@ def test_prefixed_csrf(self):
         assert response.status_code == 200
 
     def test_invalid_secure_csrf(self):
-        response = self.client.get("/", base_url='https://localhost/')
-        csrf_token = get_csrf_token(response.data)
+        with self.client:
+            self.client.get('/', base_url='https://localhost/')
+            csrf_token = request.csrf_token
 
         response = self.client.post(
             "/",
@@ -161,8 +156,10 @@ def test_invalid_secure_csrf(self):
         assert b'not match' in response.data
 
     def test_valid_secure_csrf(self):
-        response = self.client.get("/", base_url='https://localhost/')
-        csrf_token = get_csrf_token(response.data)
+        with self.client:
+            self.client.get('/', base_url='https://localhost/')
+            csrf_token = request.csrf_token
+
         response = self.client.post(
             "/",
             data={"name": "danny"},
@@ -177,8 +174,9 @@ def test_valid_secure_csrf(self):
         assert response.status_code == 200
 
     def test_valid_csrf_method(self):
-        response = self.client.get("/")
-        csrf_token = get_csrf_token(response.data)
+        with self.client:
+            self.client.get('/')
+            csrf_token = request.csrf_token
 
         response = self.client.post("/csrf-protect-method", data={
             "csrf_token": csrf_token
@@ -189,17 +187,19 @@ def test_invalid_csrf_method(self):
         response = self.client.post("/csrf-protect-method", data={"name": "danny"})
         assert response.status_code == 400
 
-        @self.csrf.error_handler
-        def invalid(reason):
-            return reason
+        @self.app.errorhandler(CsrfError)
+        def handle_csrf_error(e):
+            return e, 200
 
         response = self.client.post("/", data={"name": "danny"})
         assert response.status_code == 200
         assert b'token missing' in response.data
 
     def test_empty_csrf_headers(self):
-        response = self.client.get("/", base_url='https://localhost/')
-        csrf_token = get_csrf_token(response.data)
+        with self.client:
+            self.client.get('/', base_url='https://localhost/')
+            csrf_token = request.csrf_token
+
         self.app.config['WTF_CSRF_HEADERS'] = list()
         response = self.client.post(
             "/",
@@ -215,8 +215,10 @@ def test_empty_csrf_headers(self):
         assert response.status_code == 400
 
     def test_custom_csrf_headers(self):
-        response = self.client.get("/", base_url='https://localhost/')
-        csrf_token = get_csrf_token(response.data)
+        with self.client:
+            self.client.get('/', base_url='https://localhost/')
+            csrf_token = request.csrf_token
+
         self.app.config['WTF_CSRF_HEADERS'] = ['X-XSRF-TOKEN']
         response = self.client.post(
             "/",
@@ -239,9 +241,10 @@ def test_testing(self):
         self.app.testing = True
         self.client.post("/", data={"name": "danny"})
 
-    def test_csrf_exempt(self):
-        response = self.client.get("/csrf-exempt")
-        csrf_token = get_csrf_token(response.data)
+    def test_csrf_exempt_view_with_form(self):
+        with self.client:
+            self.client.get('/', base_url='https://localhost/')
+            csrf_token = request.csrf_token
 
         response = self.client.post("/csrf-exempt", data={
             "name": "danny",
@@ -257,7 +260,7 @@ def test_validate_csrf(self):
 
     def test_validate_not_expiring_csrf(self):
         with self.app.test_request_context():
-            csrf_token = generate_csrf(time_limit=False)
+            csrf_token = generate_csrf()
             assert validate_csrf(csrf_token, time_limit=False)
 
     def test_csrf_token_helper(self):
@@ -265,8 +268,9 @@ def test_csrf_token_helper(self):
         def withtoken():
             return render_template("csrf.html")
 
-        response = self.client.get('/token')
-        assert b'#' in response.data
+        with self.client:
+            response = self.client.get('/token')
+            assert re.search(br'token: ([0-9a-zA-Z\-._]+)', response.data)
 
     def test_csrf_blueprint(self):
         response = self.client.post('/bar/foo')
@@ -281,8 +285,9 @@ def test_csrf_token_macro(self):
         def withtoken():
             return render_template("import_csrf.html")
 
-        response = self.client.get('/token')
-        assert b'#' in response.data
+        with self.client:
+            response = self.client.get('/token')
+            assert request.csrf_token in response.data.decode('utf8')
 
     def test_csrf_custom_token_key(self):
         with self.app.test_request_context():
@@ -296,16 +301,31 @@ def test_csrf_custom_token_key(self):
             # However, the custom key can validate as well
             assert validate_csrf(custom_csrf_token, token_key='oauth_state')
 
-    def test_csrf_url_safe(self):
-        with self.app.test_request_context():
-            # Generate a normal and URL safe CSRF token
-            default_csrf_token = generate_csrf()
-            url_safe_csrf_token = generate_csrf(url_safe=True)
+    def test_old_error_handler(self):
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter('always', FlaskWTFDeprecationWarning)
+
+            @self.csrf.error_handler
+            def handle_csrf_error(reason):
+                return 'caught csrf return'
+
+            self.assertEqual(len(w), 1)
+            assert issubclass(w[0].category, FlaskWTFDeprecationWarning)
+            assert 'app.errorhandler(CsrfError)' in str(w[0].message)
+
+            rv = self.client.post('/', data={'name': 'david'})
+            assert b'caught csrf return' in rv.data
+
+        with warnings.catch_warnings(record=True) as w:
+            warnings.simplefilter('always', FlaskWTFDeprecationWarning)
+
+            @self.csrf.error_handler
+            def handle_csrf_error(reason):
+                abort(401, 'caught csrf abort')
 
-            # Verify they are not the same and the URL one is truly URL safe
-            assert default_csrf_token != url_safe_csrf_token
-            assert '#' not in url_safe_csrf_token
-            assert re.match(r'^[a-f0-9]+--[a-f0-9]+$', url_safe_csrf_token)
+            self.assertEqual(len(w), 1)
+            assert issubclass(w[0].category, FlaskWTFDeprecationWarning)
+            assert 'app.errorhandler(CsrfError)' in str(w[0].message)
 
-            # Verify we can validate our URL safe key
-            assert validate_csrf(url_safe_csrf_token, url_safe=True)
+            rv = self.client.post('/', data={'name': 'david'})
+            assert b'caught csrf abort' in rv.data
diff --git a/tests/test_validation.py b/tests/test_validation.py
--- a/tests/test_validation.py
+++ b/tests/test_validation.py
@@ -1,18 +1,8 @@
 from __future__ import with_statement
 
-import re
+from flask import request
 
-from .base import TestCase, MyForm, to_unicode
-
-csrf_token_input = re.compile(
-    r'name="csrf_token" type="hidden" value="([0-9a-z#A-Z-\.]*)"'
-)
-
-
-def get_csrf_token(data):
-    match = csrf_token_input.search(to_unicode(data))
-    assert match
-    return match.groups()[0]
+from .base import MyForm, TestCase, to_unicode
 
 
 class TestValidateOnSubmit(TestCase):
@@ -93,18 +83,20 @@ def test_ajax(self):
         assert response.status_code == 200
 
     def test_valid_csrf(self):
+        with self.client:
+            self.client.get('/')
+            csrf_token = request.csrf_token
 
-        response = self.client.get("/")
-        csrf_token = get_csrf_token(response.data)
-
-        response = self.client.post("/", data={"name": "danny",
-                                               "csrf_token": csrf_token})
+        response = self.client.post('/', data={
+            'name': 'danny',
+            'csrf_token': csrf_token
+        })
         assert b'DANNY' in response.data
 
     def test_double_csrf(self):
-
-        response = self.client.get("/")
-        csrf_token = get_csrf_token(response.data)
+        with self.client:
+            self.client.get('/')
+            csrf_token = request.csrf_token
 
         response = self.client.post("/two_forms/", data={
             "name": "danny",
@@ -114,6 +106,4 @@ def test_double_csrf(self):
 
     def test_valid_csrf_data(self):
         with self.app.test_request_context():
-            form = MyForm()
-            csrf_token = get_csrf_token(form.csrf_token())
-            assert form.validate_csrf_data(csrf_token)
+            assert MyForm().validate_csrf_data(request.csrf_token)