dev step

main.py

@@ -5,22 +5,23 @@ import os, json
 import time
 from hashlib import sha256
 from modules import model, oauth, security, log_module, llm, chat_functions, settings
-from modules.model import User
+from modules.model import User, Session
 from fastapi.staticfiles import StaticFiles
 from fastapi.templating import Jinja2Templates
 from datetime import datetime 
+from typing import Annotated
 
+####################### APP SETUP ########################
 app = FastAPI(docs_url=None, redoc_url=None)
-
 app.mount("/static", StaticFiles(directory="static"), name="static")
 templates = Jinja2Templates(directory="templates")
 
 fecha_unix = str(int(time.time()))
 log_module.log_write("master", "iniciado", "-")
+##########################################################
 
 
-
-########################## MAIN ##########################
+######################## ROUTE: / ########################
 @app.get("/", response_class=HTMLResponse)
 async def main_page(request: Request, user: User = Depends(User.find_from_cookie)):    
     if not user.can_use("chat"):
@@ -35,80 +36,75 @@ async def main_page(request: Request, user: User = Depends(User.find_from_cookie
             "tools": tools
         })
     return response
+##########################################################
 
 
-########################## SECURITY ##########################
+
+#################### SECURITY (OAUTH) ####################
 @app.get("/login", response_class=HTMLResponse)
 async def login(request: Request):
-    ret = templates.TemplateResponse("iniciar_sesion.html", {"request": request, "redirecturi": settings.OAUTH_REDIRECT})
+    # Shows the Start "session with google" button, and deletes token cookie
+    ret = templates.TemplateResponse("login.html", {"request": request, "redirecturi": settings.OAUTH_REDIRECT})
     ret.delete_cookie("token")
     return ret
 
 @app.get("/oauth", response_class=HTMLResponse)
 async def validate_oauth(request: Request):
+    # Get the oauth get params, 
+    # look for the google validation and info,
+    # set the session cookie and save user in DB
+
+    # Extract the Get params
     params = dict(request.query_params)
+
+    # Client browser fingerprint bypassed with google oauth as "state"
     fingerprint = params["state"].split("=",1)[1]
+
+    # Get the google user info
     google_userinfo = oauth.validate_redirect(params)
+
+    # Create the user model with the google info
     user:User = model.User.find_or_create(google_userinfo, fingerprint)
+
+    # Prepare redirect response and set the session cookie
     response = RedirectResponse(url='/')
     token = user.create_cookie()
     security.set_cookie(response, key="token", value=token)
+
+    # Saves the user
     user.update_user()
+
+
     return response
+##########################################################
 
 
-########################## API ##########################    
+########################## APIs ##########################   
+
+User_find_from_data = Annotated[User, Depends(User.find_from_data)]
+Session_find_from_data = Annotated[Session, Depends(Session.find_from_data)]
 
 @app.post("/getConfigs")
-async def get_configs(request: Request, data:dict, user: User = Depends(User.find_from_cookie)): #= Depends(model.User.find_from_header)):
-    tokens = user.tokens
+async def get_configs(response: Response, user: User_find_from_data):  
     year, month = datetime.now().strftime("%y-%m").split("-")
-    mes = sum(tokens.get(year, {}).get(month, {"_":0}).values())
-    tokens = {"mes": mes, "total": sum([z for x in tokens.values() for y in x.values() for z in y.values()])}
-
-    return user.configs.model_dump() | {"tokens": tokens}
+    month = sum(user.tokens.get(year, {}).get(month, {"_":0}).values())
+    total = sum([z for x in user.tokens.values() for y in x.values() for z in y.values()])
+    tokens = {"month": month, "total": total}
+    user._session.create_cookie(response)
+    return user.configs.model_dump() | {"tokens": tokens, "challenge": user._session.challenge}
 
 @app.post("/setConfigs")
-async def set_configs(request: Request, configs: model.Configs, user: User = Depends(User.find_from_cookie)):
-    user.configs = configs
+async def set_configs(response: Response, user: User_find_from_data):
+    user.configs = user._data
     user.update_user()
-
-
-    return {"success":True}
+    user._session.create_cookie(response)
+    return {"success":True, "challenge": user._session.challenge}
     
 
 @app.post("/getToken")
-async def get_token(request: Request, response: Response, user: User = Depends(User.find_from_data)):
-
-    if user._session.hashed != security.sha256(user.gid + user._guid + data["fingerprint"]):
-        security.raise_401()
-    
-    
-    user._session = model.Session(**{
-        "gid": user.gid,
-        "guid": user._guid,
-        "fprint": data["fingerprint"],
-        "public_key": data["public_key"]
-    })
-
-    security.challenges[data["fingerprint"]] = user._session.challenge
-    security.set_cookie(response, "token", user.create_session_cookie(), {"hours": 24})
-    return {"success":True, "challenge": user._session.challenge}
-
-
-
-
-
-
-
-
-    #    data: model.Session = Depends(model.Session.validate)
-    # user.session = model.Session(**data)
-    # data.pop("fingerprint")
-    # security.validate_signature(**data)
-    # security.create_jwt_token(data)
-    # token = user.create_token()
-    pass
+async def get_token(response: Response, session: Session_find_from_data):
+    session.create_cookie(response)
+    return {"success":True, "challenge": session.challenge}
 
 
 @app.post("/chat")

modules/model.py

@@ -1,7 +1,7 @@
-import os
+import os, json
 from pymongo.mongo_client import MongoClient
 from pymongo.server_api import ServerApi
-from fastapi import Request
+from fastapi import Request, Response
 from . import log_module, security, settings
 from datetime import timezone, datetime, timedelta
 from pydantic import BaseModel, Field, PrivateAttr
@@ -59,6 +59,7 @@ class Chat(BaseModel):
 
     def new_msg(self: Self):
         return Message(role="", content="")
+
 class Session(BaseModel):
     gid: str
     fprint: str
@@ -66,7 +67,8 @@ class Session(BaseModel):
     guid: str
     public_key: str = ""
     challenge: str = str(uuid.uuid4())
-    
+    data: dict = {}
+
     def __init__(self, **kwargs):
         kwargs["guid"] = kwargs.get("guid", str(uuid.uuid4())) 
         kwargs["hashed"] = security.sha256(kwargs["guid"] + kwargs["fprint"])
@@ -79,16 +81,50 @@ class Session(BaseModel):
             security.raise_401()
         return True
     
-    def create_token(self: Self):
-        self.challenge = str(uuid.uuid4()) 
-        dump = self.model_dump()
-        dump["fprint"] = security.sha256(
-            dump["fprint"]+
-            dump.pop("challenge")
+    @classmethod
+    def find_from_data(cls:Self, request: Request, data:dict) -> Self:
+        cookie_data:dict = security.token_from_cookie(request)
+
+        if "gid" not in cookie_data or "guid" not in cookie_data:
+            log_module.logger.error("Cookie without session needed data")
+            security.raise_401()
+        
+        if not (public_key := cookie_data.get("public_key", None)): # FIX Vuln Code
+            if request.scope["path"] != "/getToken":
+                log_module.logger.error(f"User without public key saved | {json.dumps(cookie_data)}")
+                security.raise_401()
+            else:
+                log_module.logger.info(f"API public key set gor user {cookie_data['gid']}")
+                public_key = data["public_key"]
+        else:
+            security.check_challenge(data["fingerprint"], cookie_data["challenge"])
+
+        session: Self = cls(
+            gid = cookie_data["gid"],
+            fprint = data["fingerprint"],
+            guid = cookie_data["guid"],
+            public_key = public_key,
+            data = data
         )
 
-        return security.create_jwt_token(dump)
+        if session.hashed != cookie_data["fprint"]:
+            log_module.logger.error(f"Fingerprint didnt match | {json.dumps(cookie_data)}")
+            security.raise_401()
+
 
+        security.add_challenge(session.fprint, session.challenge)
+        
+        return session
+    
+    def create_cookie(self:Self, response: Response):
+        jwt = security.create_jwt_token({
+            "gid":self.gid, 
+            "guid": self.guid, 
+            "fprint": self.hashed, 
+            "public_key": self.public_key,
+            "challenge": self.challenge
+             })
+        security.set_cookie(response, "token", jwt, {"hours": 24})
     
 class User(BaseModel):
     name: str
@@ -106,49 +142,55 @@ class User(BaseModel):
 
     @classmethod
     def find_or_create(cls: Self, data: dict, fprint: str)-> Self:
-        if ( found_:= DB.user.find_one({"gid":data["gid"]}) ):
-            found:Self = cls(**found_)
-            found._session = Session(gid=found.gid, fprint=fprint)
-            log_module.logger.info(f"User logged in: {found.gid} | fp: {found._session.fprint}")
-            return found
         
-        user:Self = cls(**data)
-        DB.user.insert_one(user.model_dump())
+        found = DB.user.find_one({"gid":data["gid"]})
+
+        user:Self =  cls(**found) if found else cls(**data)
+        
+        if not found:
+            DB.user.insert_one(user.model_dump())
+        
         user._session = Session(gid=user.gid, fprint=fprint)
-        log_module.logger.info(f"User created: {user.gid} | fp: {user._session.fprint}")
-        return user
 
-    def create_session(self: Self) -> str:
-        return Session({
-            "gid": self.gid, 
-            "fprint":self._session.fprint, 
-            "guid": self._session.guid
-        })
-  
+        log_module.logger.info(f"User {'logged' if found else'created'}: {user.gid} | fp: {user._session.fprint}")
+        return user
+    
     @classmethod
     def find_from_cookie(cls:Self, request: Request) -> Self:
-        data:dict = security.token_from_cookie(request)
-        found:dict = DB.user.find_one({"gid":data["gid"]})
-
-        if found and data.get("gid", None) and data.get("guid", None):
-            user: Self = cls(**found)
-            user._session = Session(gid = data["gid"], guid = data["guid"], fprint = data["fprint"])
-            return user
+        cookie_data:dict = security.token_from_cookie(request)
         
-        security.raise_307()
+        if "gid" not in cookie_data or "guid" not in cookie_data:
+            log_module.logger.error("Cookie without needed data")
+            security.raise_307()
+
+        found:dict = DB.user.find_one({"gid":cookie_data["gid"]})
+
+        if not found:
+            log_module.logger.error("User not found on DB")
+            security.raise_307()
+
+        user: Self = cls(**found)
+        user._session = Session(
+            gid = cookie_data["gid"], 
+            guid = cookie_data["guid"], 
+            fprint = cookie_data["fprint"]
+        )
+        return user
     
     @classmethod
     def find_from_data(cls:Self, request: Request, data:dict) -> Self:
-        
-        cookie_data:dict = security.token_from_cookie(request)
-        found:dict = DB.user.find_one({"gid":data["gid"]})
+        session = Session.find_from_data(request, data) 
 
-        if found and data.get("gid", None) and data.get("guid", None):
-            user: Self = cls(**found)
-            user._session = Session(gid = data["gid"], guid = data["guid"], fprint = data["fprint"])
-            return user
-        
-        security.raise_307()
+        found:dict = DB.user.find_one({"gid":session.gid})
+
+        if not found:
+            log_module.logger.error("User not found on DB")
+            security.raise_307()
+
+        user: Self = cls(**found)
+        user._session = session
+        return user
+    
     def update_description(self: Self, message: str) -> None:
         log_module.logger.info(f"Description Updated: {self.gid}")
         DB.user.update_one(
@@ -161,6 +203,7 @@ class User(BaseModel):
         return security.can_use(self.role, activity)
     
     def update_user(self: Self, message: Message = None) -> None:
+        log_module.logger.info(f"User Updated: {self.gid}")
         DB.user.update_one({"gid": self.gid}, {"$set": self.model_dump()})
 
     @staticmethod
@@ -173,14 +216,7 @@ class User(BaseModel):
         return security.create_jwt_token({
             "gid":self.gid, 
             "guid": self._session.guid, 
-            "fprint": security.sha256(self.gid + self._session.guid + self._session.fprint) })
+            "fprint": self._session.hashed})
     
-    def create_session_cookie(self:Self):
-        return security.create_jwt_token({
-            "gid":self.gid, 
-            "guid": self._guid, 
-            "fprint": security.sha256(self.gid + self._session.guid + self._session.fprint), 
-            "public_key": self._session.public_key,
-            "challenge": self._session.challenge
-             })
+
     

modules/security.py

@@ -20,7 +20,13 @@ for key in settings.USERS:
     settings.USERS[key] = sha256(password.encode('UTF-8')).hexdigest()
 
 
+def add_challenge(fprint:str, challenge:str):
+    challenges[fprint] = challenge
 
+def check_challenge(fprint:str, challenge:str):
+    if challenges.pop(fprint, None) == challenge:
+        return True
+    raise_401()
 
 def create_jwt_token(data, maxlife=settings.JWT_EXPIRATION_TIME_MINUTES_API):
     expire = datetime.utcnow() + timedelta(minutes=maxlife)    

static/js/chatHandler.js

@@ -44,13 +44,7 @@ class ChatGPT{
 
     }
 
-    obtenerToken(challenge){
-        let data = {
-            fingerprint: this.secHand.fingerprint,
-            public_key: this.secHand.publicKey,
-            operation: "challenge"
-        }
-
+    obtenerToken(){
         $.ajax({
             method: "POST",
             url: "/getToken",
@@ -58,7 +52,10 @@ class ChatGPT{
                 "Autorization": "Bearer " + this.token,
                 'Content-Type': 'application/json',
             },
-            data: JSON.stringify(data),
+            data: JSON.stringify({
+                fingerprint: this.secHand.fingerprint,
+                public_key: this.secHand.publicKey
+            }),
             timeout: 5000,
             dataType: "json"
         }).done((data) => {
@@ -84,7 +81,8 @@ class ChatGPT{
                 'Content-Type': 'application/json',
             },
             data:JSON.stringify({
-                "challenge": this.challenge
+                challenge: this.challenge,
+                fingerprint: this.secHand.fingerprint
             }),
             timeout: 5000,
             dataType: "json"
@@ -114,7 +112,7 @@ class ChatGPT{
             $(".range select").on("change", function(){
                 self.config[this.id] = this.value
             })
-            $("#tokensUsedMonth").text(data.tokens.mes)
+            $("#tokensUsedMonth").text(data.tokens.month)
             $("#tokensUsedTotal").text(data.tokens.total)
 
 
@@ -171,9 +169,16 @@ class ChatGPT{
                 "Autorization": "Bearer " + this.token,
                 'Content-Type': 'application/json',
             },
-            data: JSON.stringify(this.config),
+            data: JSON.stringify(
+                {...this.config,
+                fingerprint: this.secHand.fingerprint,
+                challenge: this.challenge,
+            }),
             timeout: 5000,
             dataType: "json"
+        }).done((data) => {
+            this.secHand.sign(data.challenge).then((result) => {this.challenge = result})
+            this.token = data.token
         })
     }