wifi-with-me/wifiwithme/apps/contribmap/tokens.py

import json

from django.core.signing import TimestampSigner, BadSignature
from django.utils.http import urlsafe_base64_decode, urlsafe_base64_encode
from django.utils.encoding import DjangoUnicodeDecodeError

from .models import Contrib


class URLTokenManager(TimestampSigner):
    """ Handle signed json data as URL-safe strings

    This class has two responsibilities:
    - sign/unsign
    - pack/unpack JSON to base64
    """

    def sign(self, payload, *args, **kwargs):
        """
        :param payload: the data to be embeded into the token
        :type data: dict
        """
        return urlsafe_base64_encode(
            super().sign(
                json.dumps(payload), *args, **kwargs
            ).encode('utf-8')).decode('ascii')

    def unsign(self, encoded_payload, *args, **kwargs):
        decoded_payload = urlsafe_base64_decode(encoded_payload)
        unsigned = super().unsign(decoded_payload)
        return json.loads(unsigned)


class TokenError(Exception):
    pass


class ContribTokenManager:
    """Produce and use signed URL tokens for account-less contrib management
    """
    SCOPE = 'contrib/manage'

    def __init__(self):
        self.manager = URLTokenManager()

    def mk_token(self, contrib):
        """ Generate a signed contrib management token

        Valid for a given contrib, and for a limited time.

        :type contrib: Contrib
        :rtype str:
        """
        return self.manager.sign({'scope': self.SCOPE, 'pk': contrib.pk})

    def get_instance_if_allowed(self, encoded_token, pk=None):
        """Return a contrib if the token grants authz for that Contrib pk

        Raise a TokenError if not authorized.

        :param pk: the contrib pk (optional, if you want to check that the
          instance is the right one)
        :param encoded_token: the encoded token, from ``mk_token()``:
        :return: a Contrib instance or None, if the object does not exist
        :rtype Contrib:
        """
        try:
            data = self.manager.unsign(encoded_token)
        except BadSignature:
            raise TokenError('Invalid token signature')

        except (DjangoUnicodeDecodeError, UnicodeDecodeError):
            raise TokenError('This is not a well-formed token')

        if (pk is not None) and (data['pk'] != pk):
            raise TokenError('Token is not valid for id {}'.format(pk))
        else:
            try:
                return Contrib.objects.get(pk=data['pk'])
            except Contrib.DoesNotExist:
                return None