From 2964b726352b4fd3d85751ea6a3c94d78d98d13b Mon Sep 17 00:00:00 2001
From: Stepan Vladovskiy <stepan.vladovskiy@gmail.com>
Date: Mon, 6 May 2024 05:24:54 -0300
Subject: [PATCH] feat: auth.py total recomp

---
 auth.py        | 34 ++++++++++++----------------------
 main.py        |  2 +-
 pyproject.toml |  1 +
 3 files changed, 14 insertions(+), 23 deletions(-)

diff --git a/auth.py b/auth.py
index 30e3d09..cae595a 100644
--- a/auth.py
+++ b/auth.py
@@ -1,22 +1,18 @@
 from functools import wraps
+from starlette.responses import JSONResponse
 import aiohttp
-from aiohttp import web
 
 AUTH_URL = 'https://auth.discours.io'
 
-
 async def check_auth(req):
     token = req.headers.get("Authorization")
-    headers = {"Authorization": token, "Content-Type": "application/json"}  # "Bearer " + removed
+    headers = {"Authorization": token, "Content-Type": "application/json"}
+
     print(f"[services.auth] checking auth token: {token}")
 
-    query_name = "session"
-    query_type = "query"
-    operation = "GetUserId"
-
     gql = {
-        "query": query_type + " " + operation + " { " + query_name + " { user { id } } }",
-        "operationName": operation,
+        "query": "query GetUserId { session { user { id } } }",
+        "operationName": "GetUserId",
         "variables": None,
     }
 
@@ -27,25 +23,19 @@ async def check_auth(req):
                 return False, None
             r = await response.json()
             if r:
-                user_id = r.get("data", {}).get(query_name, {}).get("user", {}).get("id", None)
+                user_id = r.get("data", {}).get("session", {}).get("user", {}).get("id", None)
                 is_authenticated = user_id is not None
                 return is_authenticated, user_id
     return False, None
 
-
 def login_required(f):
     @wraps(f)
-    async def decorated_function(*args, **kwargs):
-        info = args[1]
-        context = info.context
-        req = context.get("request")
-        is_authenticated, user_id = await check_auth(req)
+    async def decorated_function(request, *args, **kwargs):
+        is_authenticated, user_id = await check_auth(request)
         if not is_authenticated:
-            raise web.HTTPUnauthorized(text="You are not logged in")  # Return HTTP 401 Unauthorized
-        else:
-            context["user_id"] = user_id
-
-        # If the user is authenticated, execute the resolver
-        return await f(*args, **kwargs)
+            return JSONResponse({'error': 'Unauthorized'}, status_code=401)
 
+        # Make user_id available to the route handler, if needed
+        request.state.user_id = user_id
+        return await f(request, *args, **kwargs)
     return decorated_function
diff --git a/main.py b/main.py
index bcfe9a9..37024c7 100644
--- a/main.py
+++ b/main.py
@@ -18,7 +18,7 @@ STORJ_END_POINT = os.environ.get('STORJ_END_POINT')
 STORJ_BUCKET_NAME = os.environ.get('STORJ_BUCKET_NAME')
 CDN_DOMAIN = os.environ.get('CDN_DOMAIN')
 
-# @check_auth
+@check_auth
 async def upload_handler(request: Request):
     logging.debug("Received upload request")
     form = await request.form()
diff --git a/pyproject.toml b/pyproject.toml
index b81f107..8d5ea3a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -12,6 +12,7 @@ aiohttp = "^3.9.1"
 uvicorn = "^0.24.0.post1"
 starlette = "^0.33.0"
 aioboto3 = "^9.0.0"
+python-multipart = "^0.0.5"
 
 [tool.poetry.dev-dependencies]
 black = "^23.10.1"