codekingpro/portable-devtools

1"""2A provided CSRF implementation which puts CSRF data in a session.3 4This can be used fairly comfortably with many `request.session` type5objects, including the Werkzeug/Flask session store, Django sessions, and6potentially other similar objects which use a dict-like API for storing7session keys.8 9The basic concept is a randomly generated value is stored in the user's10session, and an hmac-sha1 of it (along with an optional expiration time,11for extra security) is used as the value of the csrf_token. If this token12validates with the hmac of the random value + expiration time, and the13expiration time is not passed, the CSRF validation will pass.14"""15import hmac16import os17from datetime import datetime18from datetime import timedelta19from hashlib import sha120 21from ..validators import ValidationError22from .core import CSRF23 24__all__ = ("SessionCSRF",)25 26 27class SessionCSRF(CSRF):28    TIME_FORMAT = "%Y%m%d%H%M%S"29 30    def setup_form(self, form):31        self.form_meta = form.meta32        return super().setup_form(form)33 34    def generate_csrf_token(self, csrf_token_field):35        meta = self.form_meta36        if meta.csrf_secret is None:37            raise Exception(38                "must set `csrf_secret` on class Meta for SessionCSRF to work"39            )40        if meta.csrf_context is None:41            raise TypeError("Must provide a session-like object as csrf context")42 43        session = self.session44 45        if "csrf" not in session:46            session["csrf"] = sha1(os.urandom(64)).hexdigest()47 48        if self.time_limit:49            expires = (self.now() + self.time_limit).strftime(self.TIME_FORMAT)50            csrf_build = "{}{}".format(session["csrf"], expires)51        else:52            expires = ""53            csrf_build = session["csrf"]54 55        hmac_csrf = hmac.new(56            meta.csrf_secret, csrf_build.encode("utf8"), digestmod=sha157        )58        return f"{expires}##{hmac_csrf.hexdigest()}"59 60    def validate_csrf_token(self, form, field):61        meta = self.form_meta62        if not field.data or "##" not in field.data:63            raise ValidationError(field.gettext("CSRF token missing."))64 65        expires, hmac_csrf = field.data.split("##", 1)66 67        check_val = (self.session["csrf"] + expires).encode("utf8")68 69        hmac_compare = hmac.new(meta.csrf_secret, check_val, digestmod=sha1)70        if hmac_compare.hexdigest() != hmac_csrf:71            raise ValidationError(field.gettext("CSRF failed."))72 73        if self.time_limit:74            now_formatted = self.now().strftime(self.TIME_FORMAT)75            if now_formatted > expires:76                raise ValidationError(field.gettext("CSRF token expired."))77 78    def now(self):79        """80        Get the current time. Used for test mocking/overriding mainly.81        """82        return datetime.now()83 84    @property85    def time_limit(self):86        return getattr(self.form_meta, "csrf_time_limit", timedelta(minutes=30))87 88    @property89    def session(self):90        return getattr(91            self.form_meta.csrf_context, "session", self.form_meta.csrf_context92        )93