cookie.py 6.4 KB

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