codekingpro/portable-devtools

1from wtforms.fields import HiddenField2from wtforms.validators import ValidationError3 4__all__ = ("CSRFTokenField", "CSRF")5 6 7class CSRFTokenField(HiddenField):8    """9    A subclass of HiddenField designed for sending the CSRF token that is used10    for most CSRF protection schemes.11 12    Notably different from a normal field, this field always renders the13    current token regardless of the submitted value, and also will not be14    populated over to object data via populate_obj15    """16 17    current_token = None18 19    def __init__(self, *args, **kw):20        self.csrf_impl = kw.pop("csrf_impl")21        super().__init__(*args, **kw)22 23    def _value(self):24        """25        We want to always return the current token on render, regardless of26        whether a good or bad token was passed.27        """28        return self.current_token29 30    def populate_obj(self, *args):31        """32        Don't populate objects with the CSRF token33        """34        pass35 36    def pre_validate(self, form):37        """38        Handle validation of this token field.39        """40        self.csrf_impl.validate_csrf_token(form, self)41 42    def process(self, *args, **kwargs):43        super().process(*args, **kwargs)44        self.current_token = self.csrf_impl.generate_csrf_token(self)45 46 47class CSRF:48    field_class = CSRFTokenField49 50    def setup_form(self, form):51        """52        Receive the form we're attached to and set up fields.53 54        The default implementation creates a single field of55        type :attr:`field_class` with name taken from the56        ``csrf_field_name`` of the class meta.57 58        :param form:59            The form instance we're attaching to.60        :return:61            A sequence of `(field_name, unbound_field)` 2-tuples which62            are unbound fields to be added to the form.63        """64        meta = form.meta65        field_name = meta.csrf_field_name66        unbound_field = self.field_class(label="CSRF Token", csrf_impl=self)67        return [(field_name, unbound_field)]68 69    def generate_csrf_token(self, csrf_token_field):70        """71        Implementations must override this to provide a method with which one72        can get a CSRF token for this form.73 74        A CSRF token is usually a string that is generated deterministically75        based on some sort of user data, though it can be anything which you76        can validate on a subsequent request.77 78        :param csrf_token_field:79            The field which is being used for CSRF.80        :return:81            A generated CSRF string.82        """83        raise NotImplementedError()84 85    def validate_csrf_token(self, form, field):86        """87        Override this method to provide custom CSRF validation logic.88 89        The default CSRF validation logic simply checks if the recently90        generated token equals the one we received as formdata.91 92        :param form: The form which has this CSRF token.93        :param field: The CSRF token field.94        """95        if field.current_token != field.data:96            raise ValidationError(field.gettext("Invalid CSRF Token."))97