1import time
  2import zlib
  3
  4from plain.signing import (
  5    JSONSerializer,
  6    SignatureExpired,
  7    Signer,
  8    b62_decode,
  9    b62_encode,
 10    b64_decode,
 11    b64_encode,
 12)
 13
 14
 15class ExpiringSigner(Signer):
 16    """A signer with an embedded expiration (vs max age unsign)"""
 17
 18    def sign(self, value, expires_in):
 19        timestamp = b62_encode(int(time.time() + expires_in))
 20        value = f"{value}{self.sep}{timestamp}"
 21        return super().sign(value)
 22
 23    def unsign(self, value):
 24        """
 25        Retrieve original value and check the expiration hasn't passed.
 26        """
 27        result = super().unsign(value)
 28        value, timestamp = result.rsplit(self.sep, 1)
 29        timestamp = b62_decode(timestamp)
 30        if timestamp < time.time():
 31            raise SignatureExpired("Signature expired")
 32        return value
 33
 34    def sign_object(
 35        self, obj, serializer=JSONSerializer, compress=False, expires_in=None
 36    ):
 37        """
 38        Return URL-safe, hmac signed base64 compressed JSON string.
 39
 40        If compress is True (not the default), check if compressing using zlib
 41        can save some space. Prepend a '.' to signify compression. This is
 42        included in the signature, to protect against zip bombs.
 43
 44        The serializer is expected to return a bytestring.
 45        """
 46        data = serializer().dumps(obj)
 47        # Flag for if it's been compressed or not.
 48        is_compressed = False
 49
 50        if compress:
 51            # Avoid zlib dependency unless compress is being used.
 52            compressed = zlib.compress(data)
 53            if len(compressed) < (len(data) - 1):
 54                data = compressed
 55                is_compressed = True
 56        base64d = b64_encode(data).decode()
 57        if is_compressed:
 58            base64d = "." + base64d
 59        return self.sign(base64d, expires_in)
 60
 61    def unsign_object(self, signed_obj, serializer=JSONSerializer):
 62        # Signer.unsign() returns str but base64 and zlib compression operate
 63        # on bytes.
 64        base64d = self.unsign(signed_obj).encode()
 65        decompress = base64d[:1] == b"."
 66        if decompress:
 67            # It's compressed; uncompress it first.
 68            base64d = base64d[1:]
 69        data = b64_decode(base64d)
 70        if decompress:
 71            data = zlib.decompress(data)
 72        return serializer().loads(data)
 73
 74
 75def dumps(
 76    obj,
 77    key=None,
 78    salt="plain.loginlink",
 79    serializer=JSONSerializer,
 80    compress=False,
 81    expires_in=None,
 82):
 83    """
 84    Return URL-safe, hmac signed base64 compressed JSON string. If key is
 85    None, use settings.SECRET_KEY instead. The hmac algorithm is the default
 86    Signer algorithm.
 87
 88    If compress is True (not the default), check if compressing using zlib can
 89    save some space. Prepend a '.' to signify compression. This is included
 90    in the signature, to protect against zip bombs.
 91
 92    Salt can be used to namespace the hash, so that a signed string is
 93    only valid for a given namespace. Leaving this at the default
 94    value or re-using a salt value across different parts of your
 95    application without good cause is a security risk.
 96
 97    The serializer is expected to return a bytestring.
 98    """
 99    return ExpiringSigner(key=key, salt=salt).sign_object(
100        obj, serializer=serializer, compress=compress, expires_in=expires_in
101    )
102
103
104def loads(
105    s,
106    key=None,
107    salt="plain.loginlink",
108    serializer=JSONSerializer,
109    fallback_keys=None,
110):
111    """
112    Reverse of dumps(), raise BadSignature if signature fails.
113
114    The serializer is expected to accept a bytestring.
115    """
116    return ExpiringSigner(
117        key=key, salt=salt, fallback_keys=fallback_keys
118    ).unsign_object(
119        s,
120        serializer=serializer,
121    )