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