From 197617a77003d6c63e41bce379d34aad5fae49c7 Mon Sep 17 00:00:00 2001
From: mhdzumair <mhdzumair@gmail.com>
Date: Sun, 23 Feb 2025 21:11:14 +0530
Subject: [PATCH] Add support for passing b64 encoded userdata via Header API
 key without encrypted format

This change introduces API call with 'encoded_user_data' Header value. This simplifies the API call from other scraper Addons instead of encrypting and storing the encrypted userdata
---
 api/middleware.py | 11 ++++++++---
 utils/crypto.py   | 23 ++++++++++++++++++++++-
 2 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/api/middleware.py b/api/middleware.py
index cf9cf8f..8961d30 100644
--- a/api/middleware.py
+++ b/api/middleware.py
@@ -65,10 +65,15 @@ class SecureLoggingMiddleware(BaseHTTPMiddleware):
 class UserDataMiddleware(BaseHTTPMiddleware):
     async def dispatch(self, request: Request, call_next: Callable):
         endpoint = await find_route_handler(request.app, request)
-        secret_str = request.path_params.get("secret_str")
-        # Decrypt and parse the UserData from secret_str
+
+        # Decrypt and parse the UserData from secret_str or Decode from encoded_user_data header
         try:
-            user_data = await crypto_utils.decrypt_user_data(secret_str)
+            encoded_user_data = request.headers.get("encoded_user_data")
+            if encoded_user_data:
+                user_data = crypto_utils.decode_user_data(encoded_user_data)
+            else:
+                secret_str = request.path_params.get("secret_str")
+                user_data = await crypto_utils.decrypt_user_data(secret_str)
         except (ValueError, ValidationError):
             # check if the endpoint is for /streams
             if endpoint and endpoint.__name__ == "get_streams":
diff --git a/utils/crypto.py b/utils/crypto.py
index abec6db..c532597 100644
--- a/utils/crypto.py
+++ b/utils/crypto.py
@@ -14,7 +14,6 @@ from Crypto.Util.Padding import pad, unpad
 from db.config import settings
 from db.redis_database import REDIS_ASYNC_CLIENT
 from db.schemas import UserData
-from utils.runtime_const import SECRET_KEY
 
 logger = logging.getLogger(__name__)
 
@@ -153,6 +152,28 @@ class CryptoUtils:
             logger.error(f"Failed to decrypt user data: {e}")
             raise ValueError("Invalid user data")
 
+    def decode_user_data(self, encoded_user_data: str) -> UserData:
+        """Decode and decrypt user data from URL-safe string"""
+        try:
+            json_str = from_urlsafe(encoded_user_data)
+            return UserData.model_validate_json(json_str)
+        except Exception as e:
+            raise ValueError("Invalid user data")
+
+    def encode_user_data(self, user_data: UserData) -> str:
+        """Encode and encrypt user data to URL-safe string"""
+        try:
+            json_str = user_data.model_dump_json(
+                exclude_none=True,
+                exclude_defaults=True,
+                exclude_unset=True,
+                round_trip=True,
+                by_alias=True,
+            )
+            return make_urlsafe(json_str.encode("utf-8"))
+        except Exception as e:
+            raise ValueError("Failed to encode user data")
+
     async def retrieve_and_decrypt(self, storage_id: str) -> UserData:
         """Retrieve and decrypt user data from Redis"""
         if not storage_id or len(storage_id) < 37: