From 5d9b29e2ae5abfe29cbcdd3d21ab1ce3ab8d1689 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 5 Jul 2024 13:11:03 -0300 Subject: [PATCH] fix: make sure old secret keys are handled as before (#2557) * refactor: add padding function for string in auth utils * fix(auth utils): run add_padding if the secret_key is valid --- src/backend/base/langflow/services/auth/utils.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/src/backend/base/langflow/services/auth/utils.py b/src/backend/base/langflow/services/auth/utils.py index f5fba63a8..d71d9430c 100644 --- a/src/backend/base/langflow/services/auth/utils.py +++ b/src/backend/base/langflow/services/auth/utils.py @@ -332,6 +332,12 @@ 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: @@ -339,16 +345,14 @@ def ensure_valid_key(s: str) -> bytes: random.seed(s) # Generate 32 random bytes key = bytes(random.getrandbits(8) for _ in range(32)) + key = base64.urlsafe_b64encode(key) 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) + key = add_padding(s).encode() + return key def get_fernet(settings_service=Depends(get_settings_service)): - SECRET_KEY = settings_service.auth_settings.SECRET_KEY.get_secret_value() + SECRET_KEY: str = settings_service.auth_settings.SECRET_KEY.get_secret_value() valid_key = ensure_valid_key(SECRET_KEY) fernet = Fernet(valid_key) return fernet