web/lib/django/contrib/messages/storage/cookie.py
author ymh <ymh.work@gmail.com>
Tue, 25 May 2010 02:43:45 +0200
changeset 29 cc9b7e14412b
permissions -rw-r--r--
update django and lucene
Ignore whitespace changes - Everywhere: Within whitespace: At end of lines:
29
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
     1
import hmac
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
     2
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
     3
from django.conf import settings
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
     4
from django.contrib.messages import constants
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
     5
from django.contrib.messages.storage.base import BaseStorage, Message
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
     6
from django.http import CompatCookie
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
     7
from django.utils import simplejson as json
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
     8
from django.utils.hashcompat import sha_hmac
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
     9
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    10
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    11
class MessageEncoder(json.JSONEncoder):
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    12
    """
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    13
    Compactly serializes instances of the ``Message`` class as JSON.
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    14
    """
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    15
    message_key = '__json_message'
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    16
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    17
    def default(self, obj):
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    18
        if isinstance(obj, Message):
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    19
            message = [self.message_key, obj.level, obj.message]
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    20
            if obj.extra_tags:
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    21
                message.append(obj.extra_tags)
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    22
            return message
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    23
        return super(MessageEncoder, self).default(obj)
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    24
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    25
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    26
class MessageDecoder(json.JSONDecoder):
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    27
    """
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    28
    Decodes JSON that includes serialized ``Message`` instances.
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    29
    """
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    30
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    31
    def process_messages(self, obj):
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    32
        if isinstance(obj, list) and obj:
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    33
            if obj[0] == MessageEncoder.message_key:
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    34
                return Message(*obj[1:])
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    35
            return [self.process_messages(item) for item in obj]
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    36
        if isinstance(obj, dict):
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    37
            return dict([(key, self.process_messages(value))
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    38
                         for key, value in obj.iteritems()])
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    39
        return obj
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    40
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    41
    def decode(self, s, **kwargs):
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    42
        decoded = super(MessageDecoder, self).decode(s, **kwargs)
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    43
        return self.process_messages(decoded)
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    44
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    45
class CookieStorage(BaseStorage):
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    46
    """
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    47
    Stores messages in a cookie.
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    48
    """
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    49
    cookie_name = 'messages'
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    50
    # We should be able to store 4K in a cookie, but Internet Explorer
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    51
    # imposes 4K as the *total* limit for a domain.  To allow other
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    52
    # cookies, we go for 3/4 of 4K.
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    53
    max_cookie_size = 3072
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    54
    not_finished = '__messagesnotfinished__'
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    55
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    56
    def _get(self, *args, **kwargs):
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    57
        """
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    58
        Retrieves a list of messages from the messages cookie.  If the
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    59
        not_finished sentinel value is found at the end of the message list,
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    60
        remove it and return a result indicating that not all messages were
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    61
        retrieved by this storage.
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    62
        """
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    63
        data = self.request.COOKIES.get(self.cookie_name)
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    64
        messages = self._decode(data)
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    65
        all_retrieved = not (messages and messages[-1] == self.not_finished)
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    66
        if messages and not all_retrieved:
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    67
            # remove the sentinel value
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    68
            messages.pop()
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    69
        return messages, all_retrieved
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    70
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    71
    def _update_cookie(self, encoded_data, response):
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    72
        """
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    73
        Either sets the cookie with the encoded data if there is any data to
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    74
        store, or deletes the cookie.
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    75
        """
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    76
        if encoded_data:
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    77
            response.set_cookie(self.cookie_name, encoded_data)
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    78
        else:
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    79
            response.delete_cookie(self.cookie_name)
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    80
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    81
    def _store(self, messages, response, remove_oldest=True, *args, **kwargs):
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    82
        """
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    83
        Stores the messages to a cookie, returning a list of any messages which
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    84
        could not be stored.
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    85
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    86
        If the encoded data is larger than ``max_cookie_size``, removes
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    87
        messages until the data fits (these are the messages which are
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    88
        returned), and add the not_finished sentinel value to indicate as much.
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    89
        """
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    90
        unstored_messages = []
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    91
        encoded_data = self._encode(messages)
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    92
        if self.max_cookie_size:
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    93
            # data is going to be stored eventually by CompatCookie, which
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    94
            # adds it's own overhead, which we must account for.
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    95
            cookie = CompatCookie() # create outside the loop
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    96
            def stored_length(val):
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    97
                return len(cookie.value_encode(val)[1])
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    98
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
    99
            while encoded_data and stored_length(encoded_data) > self.max_cookie_size:
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   100
                if remove_oldest:
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   101
                    unstored_messages.append(messages.pop(0))
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   102
                else:
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   103
                    unstored_messages.insert(0, messages.pop())
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   104
                encoded_data = self._encode(messages + [self.not_finished],
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   105
                                            encode_empty=unstored_messages)
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   106
        self._update_cookie(encoded_data, response)
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   107
        return unstored_messages
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   108
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   109
    def _hash(self, value):
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   110
        """
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   111
        Creates an HMAC/SHA1 hash based on the value and the project setting's
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   112
        SECRET_KEY, modified to make it unique for the present purpose.
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   113
        """
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   114
        key = 'django.contrib.messages' + settings.SECRET_KEY
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   115
        return hmac.new(key, value, sha_hmac).hexdigest()
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   116
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   117
    def _encode(self, messages, encode_empty=False):
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   118
        """
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   119
        Returns an encoded version of the messages list which can be stored as
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   120
        plain text.
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   121
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   122
        Since the data will be retrieved from the client-side, the encoded data
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   123
        also contains a hash to ensure that the data was not tampered with.
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   124
        """
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   125
        if messages or encode_empty:
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   126
            encoder = MessageEncoder(separators=(',', ':'))
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   127
            value = encoder.encode(messages)
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   128
            return '%s$%s' % (self._hash(value), value)
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   129
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   130
    def _decode(self, data):
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   131
        """
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   132
        Safely decodes a encoded text stream back into a list of messages.
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   133
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   134
        If the encoded text stream contained an invalid hash or was in an
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   135
        invalid format, ``None`` is returned.
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   136
        """
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   137
        if not data:
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   138
            return None
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   139
        bits = data.split('$', 1)
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   140
        if len(bits) == 2:
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   141
            hash, value = bits
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   142
            if hash == self._hash(value):
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   143
                try:
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   144
                    # If we get here (and the JSON decode works), everything is
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   145
                    # good. In any other case, drop back and return None.
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   146
                    return json.loads(value, cls=MessageDecoder)
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   147
                except ValueError:
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   148
                    pass
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   149
        # Mark the data as used (so it gets removed) since something was wrong
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   150
        # with the data.
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   151
        self.used = True
cc9b7e14412b update django and lucene
ymh <ymh.work@gmail.com>
parents:
diff changeset
   152
        return None