authentication ok

.env

+secret=ASLKDHSADKJFLKSJDFKJSDLFKSJ
+algorithm=HS256
\ No newline at end of file

core/auth.py

+import time
+from typing import Dict
+import jwt
+from passlib.hash import pbkdf2_sha256
+from decouple import config
+
+
+JWT_SECRET = config("secret")
+JWT_ALGORITHM = config("algorithm")
+
+
+def token_response(token: str):
+    return {
+        "access_token": token
+    }
+
+def signJWT(user_id: str) -> Dict[str, str]:
+    payload = {
+        "user_id": user_id,
+        "expires": time.time() + 60000
+    }
+    token = jwt.encode(payload, JWT_SECRET, algorithm=JWT_ALGORITHM)
+
+    return token_response(token)
+
+def decodeJWT(token: str) -> dict:
+    try:
+        decoded_token = jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALGORITHM])
+        return decoded_token if decoded_token["expires"] >= time.time() else None
+    except:
+        return {}
+
+def hash_pass(password: str) -> str:
+    return pbkdf2_sha256.hash(password)
\ No newline at end of file

core/bearer.py

+from fastapi import Request, HTTPException, Depends
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+
+from .auth import decodeJWT
+
+
+class JWTBearer(HTTPBearer):
+    username: str = None
+
+    def __init__(self, auto_error: bool = True):
+        super(JWTBearer, self).__init__(auto_error=auto_error)
+
+    async def __call__(self, request: Request):
+        credentials: HTTPAuthorizationCredentials = await super(JWTBearer, self).__call__(request)
+        if credentials:
+            if not credentials.scheme == "Bearer":
+                raise HTTPException(status_code=403, detail="Invalid authentication scheme.")
+            cred = self.verify_jwt(credentials.credentials)
+            if not cred:
+                raise HTTPException(status_code=403, detail="Invalid token or expired token.")
+            self.username = cred['user_id'] 
+            return credentials.credentials
+        else:
+            raise HTTPException(status_code=403, detail="Invalid authorization code.")
+
+    def verify_jwt(self, jwtoken: str) -> bool:
+        isTokenValid: bool = False
+
+        try:
+            payload = decodeJWT(jwtoken)
+        except:
+            payload = None
+        if payload:
+            return payload
+        return isTokenValid
+    
+async def get_current_user(token: str = Depends(JWTBearer())) -> str:
+    payload = decodeJWT(token)
+    return payload['user_id'] 
\ No newline at end of file

models.py → core/models.py

-from pydantic import BaseModel
 from typing import Union
+from pydantic import BaseModel, Field, EmailStr
 
 class User(BaseModel):
     username: str
-    token: Union[str, None] = None
+    email: EmailStr
+    password: str 
 
 class DataCenter(BaseModel):
     name: str
 
 class Token(BaseModel):
     datacenter: str
-    token: str
\ No newline at end of file
+    token: str
+
+class UserLoginSchema(BaseModel):
+    username: EmailStr = Field(...)
+    password: str = Field(...)
+
+    class Config:
+        schema_extra = {
+            "example": {
+                "email": "abdulazeez@x.com",
+                "password": "weakpassword"
+            }
+        }
\ No newline at end of file

main.py

-from fastapi import FastAPI, Header
+from fastapi import FastAPI, Response, Body, Depends
 from fastapi.responses import ORJSONResponse
 import json
+from fastapi import HTTPException
 from typing import Union, List
 import requests
 from sredis.sredis import * 
+from sredis.models import * 
 import logging
-from models import User, DataCenter, Token
+from core.models import User, DataCenter, Token
 from uuid import uuid4
+from core.auth import signJWT, hash_pass
+from core.bearer import get_current_user
 
 logging.config.fileConfig('logging.conf', disable_existing_loggers=False)
 
 # get root logger
 logger = logging.getLogger(__name__)
+requests.packages.urllib3.disable_warnings() 
 
 app = FastAPI()
 add_datacenter("https://kappa1.fured.cloud.bme.hu")
 add_datacenter("https://kappa2.fured.cloud.bme.hu")
 add_datacenter("https://kappa3.fured.cloud.bme.hu")
 
+
+@app.post("/user/signup", tags=["user"])
+async def create_user(user: User = Body(...)):
+    user = PUser(
+        username=user.username,
+        email=user.email,
+        password=hash_pass(user.password)
+    )
+    user.save()
+    return signJWT(user.username)
+
+@app.post("/user/login", tags=["user"])
+async def user_login(user: UserLoginSchema = Body(...)):
+    if check_user(user):
+        return signJWT(user.username)
+    return {
+        "error": "Wrong login details!"
+    }
+
 def _proxy_datacenters(serverpath: str, username, method="GET", balancer_fun = rr_get):
     server = balancer_fun()
     token = get_datacenter_token(username, server)
@@ -31,7 +55,8 @@ def _proxy_datacenters(serverpath: str, username, method="GET", balancer_fun = r
             'Authorization': token
         }
     )
-    
+    if t_resp.status_code / 100 != 2:
+        raise HTTPException(status_code=t_resp.status_code, detail="Remote server error")
     response = ORJSONResponse(
         json.loads(t_resp.content), 
         status_code=t_resp.status_code
@@ -40,53 +65,31 @@ def _proxy_datacenters(serverpath: str, username, method="GET", balancer_fun = r
 
 @app.get("/lb/{server_path:path}")
 def proxy(
-        Authorization: Union[str, None] = Header(default=None), 
-        server_path: str = "/"
+        server_path: str = "/",
+        username = Depends(get_current_user)
     ):
-    username = get_user(str(Authorization))
-    if not username:
-        return Response(status_code=401)
-    
     return _proxy_datacenters(server_path, username)
 
 @app.post("/lb/{server_path:path}")
 def proxy(
-        Authorization: Union[str, None] = Header(default=None),
-        server_path: str = "/"
+        server_path: str = "/",
+        username = Depends(get_current_user)
     ):
-    username = get_user(str(Authorization))
-    if not username:
-        return Response(status_code=401)
     return _proxy_datacenters(server_path, username, method="POST")
 
 @app.post("/add_datacenter/")
 def create_datacenter(
-        Authorization: Union[str, None] = Header(default=None),
-        dc: DataCenter = None
+        dc: DataCenter = None,
+        username = Depends(get_current_user)
     ):
-    username = get_user(str(Authorization))
-    if not username:
-        return Response(status_code=401)
     add_datacenter(dc.name)
     return Response(status_code=201)
 
-
-@app.post("/create_user/")
-def create_user(
-        user: User = None
-    ):
-    user.token = str(uuid4())
-    new_user(user.username, user.token)
-    return user
-
 @app.post("/set_tokens/")
 def set_tokens(
-        Authorization: Union[str, None] = Header(default=None),
-        tokens: List[Token] = None
+        tokens: List[Token] = None,
+        username = Depends(get_current_user)
     ):
-    username = get_user(str(Authorization))
-    if not username:
-        return Response(status_code=401)
     for token in tokens:
         set_token(username, str(token.datacenter), str(token.token))
     return tokens

pyproject.toml

@@ -13,6 +13,11 @@ uvicorn = "^0.20.0"
 requests = "^2.28.2"
 redis = {extras = ["hiredis"], version = "^4.5.1"}
 orjson = "^3.8.7"
+pyjwt = "^2.6.0"
+python-decouple = "^3.8"
+redis-om = "^0.1.2"
+pydantic = {extras = ["email"], version = "^1.10.6"}
+passlib = "^1.7.4"
 
 [tool.poe.tasks.start]
 shell = "poetry run uvicorn main:app --reload --port 6973 --host 0.0.0.0"

sredis/models.py

+import datetime
+from typing import Optional
+from redis_om import HashModel
+from pydantic import EmailStr
+
+class PUser(HashModel):
+    username: str
+    email: EmailStr
+    password: str 

sredis/sredis.py

 import redis
+from core.auth import hash_pass
+from core.models import UserLoginSchema
+from .models import PUser
+from passlib.hash import pbkdf2_sha256
 
 r = redis.Redis(host='localhost', port=6379, db=0, decode_responses=True)
 r.set("datacenters_cnt", 1)
@@ -21,16 +25,15 @@ def rr_get():
         r.incr("roundrobin_cnt")
     return str(r.hget("datacenters_hash", rr))
 
-def new_user(username: str, token: str):
-    if not r.hexists("users", token):
-        r.hset("users", token, username)
-
-def get_user(token: str):
-    user = str(r.hget("users", token))
-    return user
-
 def set_token(username: str, datacenter: str, token: str):
+    print(f"tokens:{username}" + datacenter)
     r.hset(f"tokens:{username}", datacenter, token)
 
 def get_datacenter_token(username: str, datacenter: str):
-    return str(r.hget(f"tokens:{username}", datacenter))
\ No newline at end of file
+    return str(r.hget(f"tokens:{username}", datacenter))
+
+def check_user(data: UserLoginSchema):
+    user = PUser.find((PUser.username == data.username)).all()
+    if pbkdf2_sha256.verify(data.password, user[0].password):
+        return user[0] 
+    return False
\ No newline at end of file