Use URLSafeTimedSerializer for confirmation token

3 files changed, 34 insertions, 90 deletions

diff --git a/accounts/utils/__init__.py b/accounts/utils/__init__.py
index 6698734..4529796 100644
--- a/accounts/utils/__init__.py
+++ b/accounts/utils/__init__.py
@@ -1,19 +1,14 @@
 # -*- coding: utf-8 -*-
-import hmac
 import importlib
-import pickle
 import re
-import struct
-from base64 import urlsafe_b64encode, urlsafe_b64decode
 from functools import wraps
 from flask import current_app, flash, g, redirect, render_template, request, session
 from flask import url_for as flask_url_for
-from hashlib import sha1
-from itertools import izip
-from time import time
 from werkzeug.exceptions import Forbidden
 from wtforms.validators import Regexp, ValidationError
 
+from .confirmation import Confirmation
+
 
 _username_re = re.compile(r'^[a-zA-Z][a-zA-Z0-9-]{1,15}$')
 _username_exclude_re = re.compile(r'^(admin|root)')
@@ -85,83 +80,6 @@ def logout_user():
     g.user = None
 
 
-def make_confirmation(realm, data):
-    """
-    Create a confirmation token e.g. for confirmation mails.
-
-    Expects as input a realm to distinguish data for several applications and
-    some data (pickle-able).
-    """
-    key = '\0'.join((current_app.config['SECRET_KEY'], realm))
-    payload = ''.join((struct.pack('>i', time()), pickle.dumps(data)))
-    mac = hmac.new(key, payload, sha1)
-    return urlsafe_b64encode(''.join((mac.digest(), payload))).rstrip('=')
-
-class ConfirmationInvalid(ValueError):
-    """Raised by `verify_confirmation` on invalid input data"""
-
-class ConfirmationTimeout(ValueError):
-    """Raised by `verify_confirmation` when the input data is too old"""
-
-def verify_confirmation(realm, token, timeout=None):
-    """
-    Verify a confirmation token created by `make_confirmation` and, if it is
-    valid, return the original data.
-    If `timeout` is given, only accept the token if it is less than `timeout`
-    seconds old.
-    """
-    key = '\0'.join((current_app.config['SECRET_KEY'], realm))
-
-    token = urlsafe_b64decode(token + '=' * (4 - len(token) % 4))
-    mac = token[:20]
-    tokentime = struct.unpack('>i', token[20:24])[0]
-    payload = token[24:]
-
-    if not constant_time_compare(mac, hmac.new(key, token[20:], sha1).digest()):
-        raise ConfirmationInvalid('MAC does not match')
-
-    if timeout is not None and time() > tokentime + timeout:
-        raise ConfirmationTimeout('Token is too old')
-
-    return pickle.loads(payload)
-
-def http_verify_confirmation(*args, **kwargs):
-    """
-    Like `verify_confirmation`, but raise HTTP exceptions with appropriate
-    messages instead of `Confirmation{Invalid,Timeout}`.
-    """
-
-    try:
-        return verify_confirmation(*args, **kwargs)
-    except ConfirmationInvalid:
-        raise Forbidden(u'Ungültiger Bestätigungslink.')
-    except ConfirmationTimeout:
-        raise Forbidden(u'Bestätigungslink ist zu alt.')
-
-
-# Shamelessly stolen from https://github.com/mitsuhiko/itsdangerous/
-# (C) 2011 by Armin Ronacher and the Django Software Foundation. 3-clause BSD.
-def constant_time_compare(val1, val2):
-    """Returns True if the two strings are equal, False otherwise.
-
-    The time taken is independent of the number of characters that match. Do
-    not use this function for anything else than comparision with known
-    length targets.
-
-    This is should be implemented in C in order to get it completely right.
-    """
-    len_eq = len(val1) == len(val2)
-    if len_eq:
-        result = 0
-        left = val1
-    else:
-        result = 1
-        left = val2
-    for x, y in izip(bytearray(left), bytearray(val2)):
-        result |= x ^ y
-    return result == 0
-
-
 def ensure_utf8(s):
     if isinstance(s, unicode):
         s = s.encode('utf8')
@@ -169,7 +87,7 @@ def ensure_utf8(s):
 
 
 def send_register_confirmation_mail(username, mail):
-    confirm_token = make_confirmation('register', (username, mail))
+    confirm_token = Confirmation('register').dumps((username, mail))
     confirm_link = url_for('default.register_complete', token=confirm_token, _external=True)
 
     body = render_template('mail/register.txt', username=username,
diff --git a/accounts/utils/confirmation.py b/accounts/utils/confirmation.py
new file mode 100644
index 0000000..3ae66fe
--- /dev/null
+++ b/accounts/utils/confirmation.py
@@ -0,0 +1,25 @@
+# -*- coding: utf-8 -*-
+from flask import current_app
+from itsdangerous import BadSignature, SignatureExpired, URLSafeTimedSerializer
+from werkzeug.exceptions import Forbidden
+
+
+class Confirmation(URLSafeTimedSerializer):
+
+    def __init__(self, realm, key=None, **kwargs):
+        if key is None:
+            key = current_app.config['SECRET_KEY']
+        super(Confirmation, self).__init__(key, salt=realm, **kwargs)
+
+    def loads_http(self, s, max_age=None, return_timestamp=False, salt=None):
+        """
+        Like `Confirmation.loads`, but raise HTTP exceptions with appropriate
+        messages instead of `BadSignature` or `SignatureExpired`.
+        """
+
+        try:
+            return self.loads(s, max_age, return_timestamp, salt)
+        except BadSignature:
+            raise Forbidden(u'Ungültiger Bestätigungslink.')
+        except SignatureExpired:
+            raise Forbidden(u'Bestätigungslink ist zu alt.')
diff --git a/accounts/views/default/__init__.py b/accounts/views/default/__init__.py
index 0074cd9..64c855f 100644
--- a/accounts/views/default/__init__.py
+++ b/accounts/views/default/__init__.py
@@ -8,6 +8,7 @@ from flask import current_app, redirect, request, g, flash, url_for
 from accounts.forms import LoginForm, RegisterForm, RegisterCompleteForm, \
     LostPasswordForm, SettingsForm
 from accounts.utils import *
+from accounts.utils.confirmation import Confirmation
 from accounts.models import Account
 
 
@@ -53,7 +54,7 @@ def register():
 @logout_required
 def register_complete(token):
     #TODO: check for double uids and mail
-    username, mail = http_verify_confirmation('register', token.encode('ascii'), timeout=3*24*60*60)
+    username, mail = Confirmation('register').loads_http(token, max_age=3*24*60*60)
 
     try:
         current_app.user_backend.get_by_uid(username)
@@ -102,7 +103,7 @@ def lost_password():
     if form.validate_on_submit():
         #TODO: make the link only usable once (e.g include a hash of the old pw)
         # atm the only thing we do is make the link valid for only little time
-        confirm_token = make_confirmation('lost_password', (form.user.uid,))
+        confirm_token = Confirmation('lost_password').dumps(form.user.uid)
         confirm_link = url_for('.lost_password_complete', token=confirm_token, _external=True)
 
         body = render_template('mail/lost_password.txt', username=form.user.uid,
@@ -124,7 +125,7 @@ def lost_password():
 @templated('lost_password_complete.html')
 @logout_required
 def lost_password_complete(token):
-    username, = http_verify_confirmation('lost_password', token.encode('ascii'), timeout=4*60*60)
+    username = Confirmation('lost_password').loads_http(token, max_age=4*60*60)
 
     form = RegisterCompleteForm(request.form)
     if form.validate_on_submit():
@@ -162,7 +163,7 @@ def settings():
 
         elif request.form.get('submit_main'):
             if form.mail.data and form.mail.data != g.user.attributes['mail']:
-                confirm_token = make_confirmation('change_mail', (g.user.uid, form.mail.data))
+                confirm_token = Confirmation('change_mail').dumps((g.user.uid, form.mail.data))
                 confirm_link = url_for('.change_mail', token=confirm_token, _external=True)
 
                 body = render_template('mail/change_mail.txt', username=g.user.uid,
@@ -210,7 +211,7 @@ def settings():
 @bp.route('/settings/change_mail/<token>')
 @login_required
 def change_mail(token):
-    username, mail = http_verify_confirmation('change_mail', token.encode('ascii'), timeout=3*24*60*60)
+    username, mail = Confirmation('change_mail').loads_http(token, max_age=3*24*60*60)
 
     if g.user.uid != username:
         raise Forbidden(u'Bitte logge dich als der Benutzer ein, dessen E-Mail-Adresse du ändern willst.')