api/security/csrf/csrf.py

import functools
import hashlib
import os

from flask import request, current_app
from itsdangerous import BadSignature, URLSafeTimedSerializer

from api.exception import CSRFError
from api.security.csrf import CSRF_SECRET_KEY
from api.security.csrf.constants import CSRF_COOKIE_NAME, SALT, CSRF_TOKEN_HEADER

CSRF_DEFAULT_MAX_AGE = 30

digest_method = hashlib.sha256
signer_kwargs = {'signer_kwargs': {'digest_method': digest_method}}

def generate_csrf_token(secret_key, salt):
    serializer = URLSafeTimedSerializer(secret_key, salt, **signer_kwargs)
    csrf_token_value = hashlib.sha256(os.urandom(64)).hexdigest()
    csrf_token = serializer.dumps(csrf_token_value)
    return csrf_token


def parse_csrf_token(secret_key, salt, token, max_age=CSRF_DEFAULT_MAX_AGE):
    serializer = URLSafeTimedSerializer(secret_key, salt, **signer_kwargs)
    return serializer.loads(token, max_age=max_age)


def validate_csrf(secret_key, csrf_cookie, csrf_header):
    try:
        csrf_cookie = parse_csrf_token(secret_key, SALT, csrf_cookie)
        csrf_header = parse_csrf_token(secret_key, SALT, csrf_header)
    except BadSignature:
        return False

    return csrf_cookie == csrf_header

# needed to only allow tests to disable csrf
def is_csrf_enabled():
    return True


def set_csrf_cookie(resp, csrf_token, max_age=CSRF_DEFAULT_MAX_AGE):
    resp.set_cookie(CSRF_COOKIE_NAME, max_age=max_age, value=csrf_token, httponly=True, secure=True, samesite="Lax")

def remove_csrf_cookie(resp):
    resp.delete_cookie(CSRF_COOKIE_NAME)

def csrf_needed(func):
    @functools.wraps(func)
    def csrf_cookie_protect(*args, **kwargs):
        if not is_csrf_enabled():
            return func(*args, **kwargs)

        csrf_cookie = request.cookies.get(CSRF_COOKIE_NAME)
        csrf_header = request.headers.get(CSRF_TOKEN_HEADER)
        secret_key = current_app.config[CSRF_SECRET_KEY]

        if not csrf_cookie:
            raise ValueError('Missing CSRF cookie')

        if not csrf_header:
            raise ValueError('Missing CSRF header in request')

        if validate_csrf(secret_key, csrf_cookie, csrf_header):
            return func(*args, **kwargs)
        else:
            raise CSRFError("provided CSRF are not valid or values do not match")

    return csrf_cookie_protect