Use URLSafeTimedSerializer for confirmation token

2 files changed, 28 insertions, 85 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.')