From 0d00df81944fb197ca027bf644f29ed0a26e9d96 Mon Sep 17 00:00:00 2001
From: Valentin CZERYBA <valcze80@gmail.com>
Date: Fri, 20 Oct 2023 22:26:48 +0200
Subject: [PATCH] auth with cookie

---
 app/dependencies/cookie.py      | 38 +++++++++++++++++++++++++++++++++
 app/dependencies/users_token.py |  4 ++--
 app/routers/token.py            |  2 +-
 3 files changed, 41 insertions(+), 3 deletions(-)
 create mode 100644 app/dependencies/cookie.py

diff --git a/app/dependencies/cookie.py b/app/dependencies/cookie.py
new file mode 100644
index 0000000..ed78336
--- /dev/null
+++ b/app/dependencies/cookie.py
@@ -0,0 +1,38 @@
+from fastapi.security import OAuth2
+from fastapi.openapi.models import OAuthFlows as OAuthFlowsModel
+from fastapi import Request
+from fastapi.security.utils import get_authorization_scheme_param
+from fastapi import HTTPException
+from fastapi import status
+from typing import Optional
+from typing import Dict
+
+
+class OAuth2PasswordBearerWithCookie(OAuth2):
+    def __init__(
+        self,
+        tokenUrl: str,
+        scheme_name: Optional[str] = None,
+        scopes: Optional[Dict[str, str]] = None,
+        auto_error: bool = True,
+    ):
+        if not scopes:
+            scopes = {}
+        flows = OAuthFlowsModel(password={"tokenUrl": tokenUrl, "scopes": scopes})
+        super().__init__(flows=flows, scheme_name=scheme_name, auto_error=auto_error)
+
+    async def __call__(self, request: Request) -> Optional[str]:
+        authorization: str = request.cookies.get("access_token")  #changed to accept access token from httpOnly Cookie
+        print("access_token is",authorization)
+
+        scheme, param = get_authorization_scheme_param(authorization)
+        if not authorization or scheme.lower() != "bearer":
+            if self.auto_error:
+                raise HTTPException(
+                    status_code=status.HTTP_401_UNAUTHORIZED,
+                    detail="Not authenticated",
+                    headers={"WWW-Authenticate": "Bearer"},
+                )
+            else:
+                return None
+        return param
\ No newline at end of file
diff --git a/app/dependencies/users_token.py b/app/dependencies/users_token.py
index 341075e..8ef3032 100644
--- a/app/dependencies/users_token.py
+++ b/app/dependencies/users_token.py
@@ -8,14 +8,14 @@ from jose import JWTError, jwt
 from passlib.context import CryptContext
 
 from ..models import users, token
-from ..dependencies import database
+from ..dependencies import database, cookie
 
 SECRET_KEY = "09d25e094faa6ca2556c818166b7a9563b93f7099f6f0f4caa6cf63b88e8d3e7"
 ALGORITHM = "HS256"
 
 pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
 
-oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
+oauth2_scheme = cookie.OAuth2PasswordBearerWithCookie(tokenUrl="token")
 
 def verify_password(plain_password, hashed_password):
     return pwd_context.verify(plain_password, hashed_password)
diff --git a/app/routers/token.py b/app/routers/token.py
index d8eae4c..02defe6 100644
--- a/app/routers/token.py
+++ b/app/routers/token.py
@@ -27,7 +27,7 @@ async def login_for_access_token(
     )
     content = {"message": "Access token generated"}
     response = JSONResponse(content=content)
-    response.set_cookie(key="jwt", value=access_token)
+    response.set_cookie(key="access_token", value="Bearer {0}".format(access_token), httponly=True)
     return response
 
 @router.get("/token",tags=["token"])