/django/contrib/messages/storage/cookie.py

https://code.google.com/p/mango-py/ · Python · 152 lines · 104 code · 13 blank · 35 comment · 19 complexity · 263427f3d586b559b8de93d9a0edb25a MD5 · raw file

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