dev test

main.py

@@ -122,7 +122,7 @@ async def get_token(response: Response, session: Session_find_from_data):
 ##########################################################
 
 @app.post("/chat")
-async def chat_async(response:Response, session: Session_find_from_data):
+async def chat_async(session: Session_find_from_data):
     body = model.Chat(messages = session.data["messages"])
     if(len(body.messages) < 1 or body.messages[-1].content==""):
         log_module.logger.warning(session.gid, "Mensaje RQ", "error, mensajes vacios")
@@ -132,7 +132,7 @@ async def chat_async(response:Response, session: Session_find_from_data):
         )
     
     response_async = llm.ejecutar(body, session)
-    session.create_cookie(response)
+    
     return StreamingResponse(llm.streamer(response_async, body, session), media_type="application/json")
 
 

modules/llm.py

@@ -95,4 +95,5 @@ async def streamer(response_async, body, session):
     session.update_usage(message)
     
     yield json.dumps({"comando":"challenge", "challenge": session.challenge} )
+    yield json.dumps({"comando":"token", "token": session.create_cookie_token() } )
     yield json.dumps({"comando":"mensaje", "mensaje": message.model_dump()} )

modules/model.py

@@ -78,7 +78,7 @@ class Session(BaseModel):
     def validate_signature(self: Self, signature:str):
         valid = security.validate_signature(self.public_key, signature, self.challenge)
         if not valid:
-            security.raise_401()
+            security.raise_401("Cannot validate Session signature")
         return True
     
     @classmethod
@@ -87,17 +87,17 @@ class Session(BaseModel):
 
         if "gid" not in cookie_data or "guid" not in cookie_data:
             log_module.logger.error("Cookie without session needed data")
-            security.raise_401()
+            security.raise_401("gid or guid not in cookie")
         
         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()
+                security.raise_401("the user must have a public key saved in token")
             else:
                 log_module.logger.info(f"API public key set for user {cookie_data['gid']}")
                 public_key = data["public_key"]
         else:
-            security.check_challenge(data["fingerprint"], cookie_data["challenge"])
+            cls.check_challenge(data["fingerprint"], cookie_data["challenge"])
 
         session: Self = cls(
             gid = cookie_data["gid"],
@@ -107,18 +107,18 @@ class Session(BaseModel):
             data = data,
             configs = Configs(**cookie_data["configs"])
         )
+        
 
         if session.hashed != cookie_data["fprint"]:
             log_module.logger.error(f"Fingerprint didnt match | {json.dumps(cookie_data)}")
-            security.raise_401()
+            security.raise_401("Fingerprint didnt match")
 
 
-        security.add_challenge(session.fprint, session.challenge)
+        session.add_challenge()
         
         return session
-    
-    def create_cookie(self:Self, response: Response):
-        jwt = security.create_jwt_token({
+    def create_cookie_token(self:Self):
+        return security.create_jwt_token({
             "gid":self.gid, 
             "guid": self.guid, 
             "fprint": self.hashed, 
@@ -126,10 +126,26 @@ class Session(BaseModel):
             "challenge": self.challenge,
             "configs": self.configs.model_dump()
              })
+    
+    def create_cookie(self:Self, response: Response):
+        jwt = self.create_cookie_token()
         security.set_cookie(response, "token", jwt, {"hours": 24})
         
     def update_usage(self: Self, message:Message):
         User.update_usage(self.gid, message)
+
+    def add_challenge(self: Self):
+        return True
+        self.challenge = str(uuid.uuid4()) 
+        DB.sess.insert_one(self.model_dump(include={"fprint", "challenge"}) )
+
+    @staticmethod 
+    def check_challenge(fprint:str, challenge: str):
+        return True
+        found = DB.sess.find_one_and_delete({"fprint":fprint})
+        if not found or found["challenge"] != challenge:
+            security.raise_401("Check challenge failed")
+            
     
 class User(BaseModel):
     name: str
@@ -166,13 +182,13 @@ class User(BaseModel):
         
         if "gid" not in cookie_data or "guid" not in cookie_data:
             log_module.logger.error("Cookie without needed data")
-            security.raise_307()
+            security.raise_307("gid or guid not in cookie")
 
         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()
+            security.raise_307("User not found on DB")
 
         user: Self = cls(**found)
         user._session = Session(
@@ -193,7 +209,7 @@ class User(BaseModel):
 
         if not found:
             log_module.logger.error("User not found on DB")
-            security.raise_307()
+            security.raise_307("User not found on DB")
 
         user: Self = cls(**found)
         user._session = session

modules/security.py

@@ -9,8 +9,6 @@ from Crypto.PublicKey import RSA
 
 import base64
 
-challenges = {}
-
 
 for key in settings.USERS:
     if key == "master": continue
@@ -18,14 +16,6 @@ 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)    
     to_encode = data | {"exp": expire}
@@ -39,18 +29,18 @@ def validate_jwt_token(token: str, usage: str = "api"):
     except Exception as e: 
         log_module.logger.error(repr(e) + " - Invalid token: " + str(token))
         if usage == "view":
-            return raise_307()
-        return raise_401()  
+            return raise_307("Failed validating jwt token")
+        return raise_401("Failed validating jwt token")  
 
 def token_from_cookie(request:Request):
     if token := request.cookies.get('token', ""):
         return validate_jwt_token(token, "view")
-    return raise_307()
+    return raise_307("Token not found")
 
 def token_from_headers(request:Request):
     if (bearer := request.headers.get("Autorization", " ").split(" ",1)) and bearer[0] == "Bearer":
         return validate_jwt_token(bearer[1])
-    return raise_401()
+    return raise_401("Bearer malformed")
     
 
 def can_use(role, activity):
@@ -59,21 +49,21 @@ def can_use(role, activity):
     }
     return role in can[activity]
 
-def raise_401():
+def raise_401(detail:str):
     headers = {"set-cookie": "token=deleted; path=/; expires=Thu, 01 Jan 1970 00:00:00 GMT",
                "Location": "/login"}
     raise HTTPException(
             status_code=status.HTTP_401_UNAUTHORIZED,
-            detail="Could not validate credentials",
+            detail=detail,
             headers=headers
         )  
     
-def raise_307():
+def raise_307(detail:str):
     headers = {"set-cookie": "token=deleted; path=/; expires=Thu, 01 Jan 1970 00:00:00 GMT",
                "Location": "/login"}
     raise HTTPException(
             status_code=status.HTTP_307_TEMPORARY_REDIRECT,
-            detail="Could not validate access",
+            detail=detail,
             headers=headers
         )  
 
@@ -85,7 +75,7 @@ def validate_signature(public_key: str, signature: str, data:str) -> bool:
         pss.new(public_key).verify(data_, signature)
         return True
     except ValueError:
-        raise_401()
+        raise_401("Signature failed")
 
 def sha256(data:str|bytes) -> str:
     data_ = data if isinstance(data, bytes) else data.encode()

modules/settings.py

@@ -16,7 +16,7 @@ OPENAI_API_KEY=os.environ.get('OPENAI_API_KEY')
 
 USERS = json.loads(str(os.getenv("USER_KEYS")).replace("\n", ""))
 
-JWT_SECRET = os.environ.get('OPENAI_API_KEY')
+JWT_SECRET = os.environ.get('JWT_SECRET')
 
 JWT_ALGORITHM = "HS256"
 JWT_EXPIRATION_TIME_MINUTES_API = 7*24*60

requirements.txt

@@ -12,6 +12,6 @@ tiktoken>=0.5.*
 pycryptodome==3.19.*
 pydantic==2.5.*
 PyJWT==2.8.*
-pymongo==4.6.*
+pymongo[srv]==4.6.*
 python-multipart>=0.0.6
 tiktoken==0.5.*

static/js/chatHandler.js

@@ -248,6 +248,10 @@ class ChatGPT{
                         case "token":
                             self.token = response.token;
                             break;
+                        // Se carga el nuevo challenge
+                        case "challenge":
+                            this.secHand.sign(response.challenge).then((result) => {self.challenge = result})
+                            break;
                         // Status del comportamiento
                         case "status":
                             ctx.respuestaStatus(response.status.mensaje, response.status.modo)
@@ -263,7 +267,7 @@ class ChatGPT{
                         //     break;
                         // Algo falló
                         default:
-                            console.log("???")
+                            console.log("???", response)
                     }
                 }
 
@@ -296,10 +300,7 @@ class ChatGPT{
                 return
             }
             return consume(response.body.getReader());
-        }).then(data => {
-            console.log("esto", data)
-        })
-        .catch(err =>{
+        }).catch(err =>{
             // Error
             console.log('Solicitud fallida', err)
             ctx.respuestaError(response)