initial commit

max/csrf.py

@@ -0,0 +1,58 @@
+from functools import wraps
+
+from flask import request, session
+
+from max.crypto import constant_time_compare, get_random_string
+from max.middlewares import BaseMiddleware
+
+
+CSRF_TOKEN_KEY = "csrf_token"
+CSRF_TOKEN_LENGTH = 32
+
+
+class CsrfMiddleware(BaseMiddleware):
+    def __call__(self, *args, **kwargs):
+        error = self.process_request()
+        if error is not None:
+            # clear csrf from session
+            if CSRF_TOKEN_KEY in session:
+                del session[CSRF_TOKEN_KEY]
+            return error, 400
+
+        return self.handler(*args, **kwargs)
+
+    def process_request(self):
+        """ If a string is returned, it mist be used as a 400er error """
+        if getattr(self.handler, "no_csrf", False):
+            return
+
+        session_has_csrf_token = CSRF_TOKEN_KEY in session
+        method_is_unsafe = request.method.lower() != "get"
+
+        if not session_has_csrf_token and method_is_unsafe:
+            return "CSRF token is missing"
+
+        elif not session_has_csrf_token and not method_is_unsafe:
+            session[CSRF_TOKEN_KEY] = get_random_string(CSRF_TOKEN_LENGTH)
+
+        elif session_has_csrf_token and method_is_unsafe:
+            return self.check_token(session[CSRF_TOKEN_KEY], request.form)
+
+    def check_token(self, session_token, formdata):
+        if CSRF_TOKEN_KEY not in formdata:
+            return "No CSRF token"
+
+        request_token = formdata[CSRF_TOKEN_KEY]
+        if len(request_token) != CSRF_TOKEN_LENGTH:
+            return "Wrong CSRF token"
+        elif not constant_time_compare(request_token, session_token):
+            return "Wrong CSRF token"
+
+
+def no_csrf(handler):
+    @wraps(handler)
+    def wrapped_handler(*args, **kwargs):
+        return handler(*args, **kwargs)
+
+    wrapped_handler.no_csrf = True
+    return wrapped_handler