Helpers for generating invitation email tokens.

invitation.py

"""
    This module contains methods for generation and validation invitation tokens.
    You should use it for implementing invitation process as following:

    1.
        # Create new inactive user
        user = User.objects.create(first_name='John', last_name='Doe', email='email@example.com', is_active=False)
        # Generate expiration token valid for two days
        token = generate_invitation_token(user.id, 60*60*24*2)
        # Send invitation email
        send_invitation_email(token)

    2.
        In finish invitation view:

        try:
            user_id = decode_invitation_token(token)
        except InvitationTokenExpiredException:
            print_expired_message()
        except InvalidInvitationTokenException:
            raise_403_exception("Invalid token")

        user = User.objects.get(pk=user_id)

        user.username = username
        user.set_password(new_password)
        user.is_active = True
        user.save()

"""

import time
import hmac
import base64
from hashlib import sha1
from settings import HMAC_KEY


class InvalidInvitationTokenException(Exception):
    pass


class InvitationTokenExpiredException(Exception):
    pass


def decode_invitation_token(token):
    """
        Decodes invitation token.

        args:
            Base64 encoded invitation token.
        returns:
            Invited user ID.
        raises:
            InvalidInvitationTokenException: when token signature or token value is invalid

            InvitationTokenExpiredException: when token expired and invitation is not more valid.
    """
    try:
        raw_token = base64.b64decode(token)
        data, signature = raw_token.split("|")
    except:
        raise InvalidInvitationTokenException()

    data_hmac = hmac.new(HMAC_KEY, data, sha1)

    if data_hmac.hexdigest() != signature:
        raise InvalidInvitationTokenException()

    action, user_id, timestamp = data.split(":")

    if action != "invitation":
        raise InvalidInvitationTokenException()

    if time.time() > float(timestamp):
        raise InvitationTokenExpiredException()

    return int(user_id)


def generate_invitation_token(user_id, max_age):
    """
        Generates invitation token.

        args:
            user_id:
                ID of newly created, but not yet activated user.
            max_age:
                Maximum age in secods until token is valid.
        returns:
            Base64 encoded invitation token
    """

    data = "invitation:{}:{}".format(str(user_id), time.time() + max_age)
    data_hmac = hmac.new(HMAC_KEY, data, sha1)
    signature = data_hmac.hexdigest()
    raw_token = "{}|{}".format(data, signature)
    return base64.b64encode(raw_token)

test_invitation.py

import time
import hmac
from hashlib import sha1
import base64
import mock

from settings import HMAC_KEY
from django.test import TestCase
from invitation import generate_invitation_token
from invitation import decode_invitation_token
from invitation import InvitationTokenExpiredException
from invitation import InvalidInvitationTokenException


class InvitationTokenTestCase(TestCase):

    def test_generate_validate_flow(self):
        token = generate_invitation_token(1, 120)
        user_id = decode_invitation_token(token)
        self.assertEquals(user_id, 1)

    def test_expired_token(self):
        token = generate_invitation_token(1, 1)
        time.sleep(2)
        self.assertRaises(InvitationTokenExpiredException, decode_invitation_token, token)

    def test_invalid_token(self):
        invalid_token = base64.b64encode("bla-bla-bla")
        self.assertRaises(InvalidInvitationTokenException, decode_invitation_token, invalid_token)

    def test_wrong_key(self):
        token = generate_invitation_token(1, 120)
        with mock.patch("consumers.invitation.HMAC_KEY", "AAAAAAAAAAA"):
            self.assertRaises(InvalidInvitationTokenException, decode_invitation_token, token)

    def test_invalid_action(self):
        data = "invalid_action:1:{}".format(time.time() + 120)
        data_hmac = hmac.new(HMAC_KEY, data, sha1)
        signature = data_hmac.hexdigest()
        raw_token = "{}|{}".format(data, signature)
        token = base64.b64encode(raw_token)

        self.assertRaises(InvalidInvitationTokenException, decode_invitation_token, token)