feat: allow cookie & JWT auth for sse/subscribe (#152)

* return cookie on login and check cookie in sse/subscribe

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* allow cookie & JWT auth for sse/subscribe

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* return token object on login to preserve compatibility

* remove unused import

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

app/main.py

@@ -2,7 +2,7 @@
 import logging
 
 from decouple import config as dconfig
-from fastapi import Depends, FastAPI, Request
+from fastapi import FastAPI, Request
 from fastapi.encoders import jsonable_encoder
 from fastapi.exceptions import HTTPException
 from fastapi_plugins import (
@@ -210,9 +210,14 @@ def _send_sse_event(id, event, data):
 @app.get(
     "/sse/subscribe",
     status_code=status.HTTP_200_OK,
-    dependencies=[Depends(JWTBearer())],
 )
 async def stream(request: Request):
+    token = request.cookies.get("access_token")
+    if not token:
+        token = request.headers.get("authorization").replace("Bearer ", "")
+        if not JWTBearer().verify_jwt(jwtoken=token):
+            raise HTTPException(401)
+
     event_source, id = sse_mgr.add_connection(request)
     new_connections.append(id)
 

app/routers/system.py

@@ -2,7 +2,7 @@
 import secrets
 
 from decouple import config
-from fastapi import APIRouter, HTTPException, Request, status
+from fastapi import APIRouter, HTTPException, Request, Response, status
 from fastapi.params import Depends
 
 from app.auth.auth_bearer import JWTBearer
@@ -40,7 +40,7 @@
     response_description="JWT token for the current session.",
     status_code=status.HTTP_200_OK,
 )
-async def login(i: LoginInput):
+async def login(i: LoginInput, response: Response):
 
     platform = ""
     try:
@@ -56,12 +56,17 @@ async def login(i: LoginInput):
             )
             data = parse_key_value_text(result)
             if data["correct"] == "1":
-                return sign_jwt()
+                tokenObj = sign_jwt()
+                token = tokenObj.get("access_token")
+                response.set_cookie("access_token", token)
+                return tokenObj
         raise HTTPException(status.HTTP_401_UNAUTHORIZED, detail="Password is wrong")
     else:
         match = secrets.compare_digest(i.password, config("login_password", cast=str))
         if match:
-            return sign_jwt()
+            token = sign_jwt().get("access_token")
+            response.set_cookie("access_token", token)
+            return token
         raise HTTPException(status.HTTP_401_UNAUTHORIZED, detail="Password is wrong")