core.py 3.1 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798
  1. from wtforms.validators import ValidationError
  2. from wtforms.fields import HiddenField
  3. __all__ = ('CSRFTokenField', 'CSRF')
  4. class CSRFTokenField(HiddenField):
  5. """
  6. A subclass of HiddenField designed for sending the CSRF token that is used
  7. for most CSRF protection schemes.
  8. Notably different from a normal field, this field always renders the
  9. current token regardless of the submitted value, and also will not be
  10. populated over to object data via populate_obj
  11. """
  12. current_token = None
  13. def __init__(self, *args, **kw):
  14. self.csrf_impl = kw.pop('csrf_impl')
  15. super(CSRFTokenField, self).__init__(*args, **kw)
  16. def _value(self):
  17. """
  18. We want to always return the current token on render, regardless of
  19. whether a good or bad token was passed.
  20. """
  21. return self.current_token
  22. def populate_obj(self, *args):
  23. """
  24. Don't populate objects with the CSRF token
  25. """
  26. pass
  27. def pre_validate(self, form):
  28. """
  29. Handle validation of this token field.
  30. """
  31. self.csrf_impl.validate_csrf_token(form, self)
  32. def process(self, *args):
  33. super(CSRFTokenField, self).process(*args)
  34. self.current_token = self.csrf_impl.generate_csrf_token(self)
  35. class CSRF(object):
  36. field_class = CSRFTokenField
  37. def setup_form(self, form):
  38. """
  39. Receive the form we're attached to and set up fields.
  40. The default implementation creates a single field of
  41. type :attr:`field_class` with name taken from the
  42. ``csrf_field_name`` of the class meta.
  43. :param form:
  44. The form instance we're attaching to.
  45. :return:
  46. A sequence of `(field_name, unbound_field)` 2-tuples which
  47. are unbound fields to be added to the form.
  48. """
  49. meta = form.meta
  50. field_name = meta.csrf_field_name
  51. unbound_field = self.field_class(
  52. label='CSRF Token',
  53. csrf_impl=self
  54. )
  55. return [(field_name, unbound_field)]
  56. def generate_csrf_token(self, csrf_token_field):
  57. """
  58. Implementations must override this to provide a method with which one
  59. can get a CSRF token for this form.
  60. A CSRF token is usually a string that is generated deterministically
  61. based on some sort of user data, though it can be anything which you
  62. can validate on a subsequent request.
  63. :param csrf_token_field:
  64. The field which is being used for CSRF.
  65. :return:
  66. A generated CSRF string.
  67. """
  68. raise NotImplementedError()
  69. def validate_csrf_token(self, form, field):
  70. """
  71. Override this method to provide custom CSRF validation logic.
  72. The default CSRF validation logic simply checks if the recently
  73. generated token equals the one we received as formdata.
  74. :param form: The form which has this CSRF token.
  75. :param field: The CSRF token field.
  76. """
  77. if field.current_token != field.data:
  78. raise ValidationError(field.gettext('Invalid CSRF Token'))