chore: Refactor authentication key generation (#2443)

Refactor the `ensure_valid_key` function in `utils.py` to improve the generation of a valid key for authentication. The function now checks if the input key is too short and generates a random key if necessary. Additionally, the key is now URL-safe base64-encoded. This change enhances the security and reliability of the authentication process.

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-07-01 13:26:35 -03:00
commit 73e38a0c73

View file

@ -1,3 +1,5 @@
import base64
import random
import warnings
from datetime import datetime, timedelta, timezone
from typing import Annotated, Coroutine, Optional, Union
@ -330,17 +332,25 @@ def authenticate_user(username: str, password: str, db: Session = Depends(get_se
return user if verify_password(password, user.password) else None
def add_padding(s):
# Calculate the number of padding characters needed
padding_needed = 4 - len(s) % 4
return s + "=" * padding_needed
def ensure_valid_key(s: str) -> bytes:
# If the key is too short, we'll use it as a seed to generate a valid key
if len(s) < 32:
# Use the input as a seed for the random number generator
random.seed(s)
# Generate 32 random bytes
key = bytes(random.getrandbits(8) for _ in range(32))
else:
# If the key is long enough, use the first 32 bytes
key = s[:32].encode()
# Ensure the key is URL-safe base64-encoded
return base64.urlsafe_b64encode(key)
def get_fernet(settings_service=Depends(get_settings_service)):
SECRET_KEY = settings_service.auth_settings.SECRET_KEY.get_secret_value()
# It's important that your secret key is 32 url-safe base64-encoded byte
padded_secret_key = add_padding(SECRET_KEY)
fernet = Fernet(padded_secret_key)
valid_key = ensure_valid_key(SECRET_KEY)
fernet = Fernet(valid_key)
return fernet