test: oauth

main.py

@@ -1,100 +1,33 @@
 from fastapi import FastAPI, Request, staticfiles, Depends, HTTPException, status
-from fastapi.responses import HTMLResponse, StreamingResponse, JSONResponse
+from fastapi.responses import HTMLResponse, StreamingResponse
 
 from fastapi.security import HTTPBasic, HTTPBasicCredentials
 import os, json
-from datetime import datetime, timedelta
 import time
-import jwt
 from openai import OpenAI
 from hashlib import sha256
 from error_map import error_table
 from logger import log_write, logger
-from oauthlib.oauth2 import WebApplicationClient
+from modules import oauth, security
 
 from pymongo.mongo_client import MongoClient
 from pymongo.server_api import ServerApi
 
-
-
-GOOGLE_CLIENT_ID = os.environ.get("GOOGLE_CSE_ID", None)
-GOOGLE_CLIENT_SECRET = os.environ.get("GOOGLE_API_KEY", None)
-GOOGLE_DISCOVERY_URL = (
-    "https://accounts.google.com/.well-known/openid-configuration"
-)
-
 MONGO_URL = os.environ.get("MONGO_URL", None)
 MONGO_PWD = os.environ.get("MONGO_PWD", None)
 MONGO_USR = os.environ.get("MONGO_USR", None)
 
-
-
 import chat_functions
 
-
-
-
 app = FastAPI()
-security = HTTPBasic()
+httpsecurity = HTTPBasic()
 model = "gpt-3.5-turbo-16k"
 app.mount("/static", staticfiles.StaticFiles(directory="static"), name="static")
 users = json.loads(str(os.getenv("USER_KEYS")).replace("\n", ""))
 
-
-for key in users:
-    if key == "master": continue
-    password = key+users[key]+users["master"]
-    users[key] = sha256(password.encode('UTF-8')).hexdigest()
-
-
 fecha_unix = str(int(time.time()))
-
-JWT_SECRET = users["master"]
-JWT_ALGORITHM = "HS256"
-JWT_EXPIRATION_TIME_MINUTES = 30
-
-
 log_write("master", "iniciado", "-")
 
-### Login stuff
-
-def authenticate_user(credentials: HTTPBasicCredentials) -> bool:
-    username = credentials.username
-    password = credentials.password
-    password = username+password+users["master"]
-    password = sha256(password.encode('UTF-8')).hexdigest()
-
-    if credentials.username not in users or password != users[credentials.username]:
-        log_write(credentials.username, "Autenticacion usuario fallida", "")
-        raise HTTPException(
-            status_code=status.HTTP_401_UNAUTHORIZED,
-            detail="Incorrect username or password",
-            headers={"WWW-Authenticate": "Basic"},
-        )
-
-    log_write(credentials.username, "Usuario autenticado", "")
-    return True
-
-
-def create_jwt_token(data):
-    to_encode = {"data": data}
-    expire = datetime.utcnow() + timedelta(minutes=JWT_EXPIRATION_TIME_MINUTES)
-    to_encode.update({"exp": expire})
-    encoded_jwt = jwt.encode(to_encode, JWT_SECRET, algorithm=JWT_ALGORITHM)
-    return encoded_jwt
-
-async def validate_token(request: Request):
-    data = {}
-    try:
-        data = await request.json()
-        token = data.pop("token")
-        payload = jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALGORITHM])
-        data["token_data"] = payload["data"]
-    except Exception as e: 
-        logger.error(repr(e) + " - " + str(data))
-        raise HTTPException(status_code=404, detail="Token inválido")    
-    return data
-
 
 
 def ejecutar_llm(data):
@@ -136,14 +69,14 @@ def manejar_funciones(chunk, response_async, data):
     return json.dumps(f_cb), ejecutar_llm(data)
 
 @app.get("/", response_class=HTMLResponse)
-async def root(request: Request, credentials: HTTPBasicCredentials = Depends(security)):
-    if authenticate_user(credentials):
-        token = create_jwt_token({"user":credentials.username})
+async def root(request: Request, credentials: HTTPBasicCredentials = Depends(httpsecurity)):
+    if security.authenticate_user(credentials):
+        token = security.create_jwt_token({"user":credentials.username})
         with open(os.path.join("static", "main.html")) as f:
             return HTMLResponse(f.read().replace("{% token %}", token).replace("{% version %}", fecha_unix))
 
 @app.post("/chat")
-async def s_chat(data = Depends(validate_token)):
+async def s_chat(data = Depends(security.validate_token)):
     messages = data.get("messages", "")
     
     if not messages:
@@ -173,7 +106,7 @@ async def s_chat(data = Depends(validate_token)):
     
     
 @app.post("/a_chat")
-async def a_chat(data = Depends(validate_token)):
+async def a_chat(data = Depends(security.validate_token)):
     messages = data.get("messages", "")
     
     
@@ -213,13 +146,13 @@ async def a_chat(data = Depends(validate_token)):
         
         mensaje = {"role": "assistant", "content":content}
         yield json.dumps({"comando": "mensaje", "mensaje": mensaje})
-        yield json.dumps({"comando": "token", "token":create_jwt_token(data["token_data"])})
+        yield json.dumps({"comando": "token", "token": security.create_jwt_token(data["token_data"])})
     return StreamingResponse(__streamer(response_async, data), media_type="application/json")
 
 
 @app.get("/read_log", response_class=HTMLResponse)
-async def read_log(request: Request, credentials: HTTPBasicCredentials = Depends(security)):
-    if sha256(credentials.username.encode()).hexdigest()=="bc1c32d709aef061bbde4fc848421cdb933e8a9f391c3a089f2861ac0772c168" and authenticate_user(credentials):
+async def read_log(request: Request, credentials: HTTPBasicCredentials = Depends(httpsecurity)):
+    if sha256(credentials.username.encode()).hexdigest()=="bc1c32d709aef061bbde4fc848421cdb933e8a9f391c3a089f2861ac0772c168" and security.authenticate_user(credentials):
         log_write(credentials.username, "Log Accesado", "")
         with open("logs/eventos.log", "r") as f:
             return HTMLResponse(f.read())
@@ -236,4 +169,22 @@ async def test(request: Request):
     except Exception as e:
         log_write("", "test", str(e))
         
+@app.get("/privacy", response_class=HTMLResponse)
+async def sesion_i(request: Request):
+    with open(os.path.join("static", "PrivacyPolicy.html")) as f:
+        return HTMLResponse(f.read())
+    
+@app.get("/login", response_class=HTMLResponse)
+async def login(request: Request):
+    with open(os.path.join("static", "iniciar_sesion.html")) as f:
+        return HTMLResponse(f.read())
+
+@app.get("/oauth", response_class=HTMLResponse)
+async def validate_oauth(request: Request):
+    params = dict(request.query_params)
+    google_userinfo = oauth.validate_redirect(params)
+    return json.dumps(google_userinfo, indent=2)
+
+
+
 

modules/oauth.py

@@ -0,0 +1,30 @@
+import requests, os
+
+
+GOOGLE_CLIENT_ID = os.environ.get("GOOGLE_CLIENT_ID", None)
+GOOGLE_CLIENT_SECRET = os.environ.get("GOOGLE_CLIENT_SECRET", None)
+OAUTH_REDIRECT = os.environ.get("OAUTH_REDIRECT", None)
+
+def validate_redirect(params):
+    data = {
+        'code': params["code"],
+        'redirect_uri': OAUTH_REDIRECT,
+        'client_id': GOOGLE_CLIENT_ID,
+        'client_secret': GOOGLE_CLIENT_SECRET,
+        'scope': params["scope"],
+        'grant_type': 'authorization_code'
+    }
+
+    google_auth_rq = requests.post("https://oauth2.googleapis.com/token", data=data)
+    google_auth = google_auth_rq.json()
+    
+    try:
+        headers = {
+            'Authorization': google_auth["token_type"]+" "+google_auth["access_token"]
+        }
+    except Exception as e:
+        print(e)
+        print(google_auth_rq.content)
+        return str(e)+" "+google_auth_rq.content
+    google_userinfo_rq = requests.get("https://www.googleapis.com/oauth2/v2/userinfo", headers=headers)
+    return google_userinfo_rq.json()

modules/security.py

@@ -0,0 +1,58 @@
+import json, os, jwt
+from hashlib import sha256
+from fastapi import HTTPException, status, Request
+from fastapi.security import HTTPBasicCredentials
+from logger import log_write, logger
+from datetime import datetime, timedelta
+
+
+
+
+users = json.loads(str(os.getenv("USER_KEYS")).replace("\n", ""))
+for key in users:
+    if key == "master": continue
+    password = key+users[key]+users["master"]
+    users[key] = sha256(password.encode('UTF-8')).hexdigest()
+
+JWT_SECRET = users["master"]
+JWT_ALGORITHM = "HS256"
+JWT_EXPIRATION_TIME_MINUTES = 30
+
+def authenticate_user(credentials: HTTPBasicCredentials) -> bool:
+    username = credentials.username
+    password = credentials.password
+    password = username+password+users["master"]
+    password = sha256(password.encode('UTF-8')).hexdigest()
+
+    if credentials.username not in users or password != users[credentials.username]:
+        log_write(credentials.username, "Autenticacion usuario fallida", "")
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Incorrect username or password",
+            headers={"WWW-Authenticate": "Basic"},
+        )
+
+    log_write(credentials.username, "Usuario autenticado", "")
+    return True
+
+def create_jwt_token(data):
+    to_encode = {"data": data}
+    expire = datetime.utcnow() + timedelta(minutes=JWT_EXPIRATION_TIME_MINUTES)
+    to_encode.update({"exp": expire})
+    encoded_jwt = jwt.encode(to_encode, JWT_SECRET, algorithm=JWT_ALGORITHM)
+    return encoded_jwt
+
+
+
+async def validate_token(request: Request):
+    data = {}
+    try:
+        data = await request.json()
+        token = data.pop("token")
+        payload = jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALGORITHM])
+        data["token_data"] = payload["data"]
+    except Exception as e: 
+        logger.error(repr(e) + " - " + str(data))
+        raise HTTPException(status_code=404, detail="Token inválido")    
+    return data
+

static/iniciar_sesion.html

@@ -1,5 +1,5 @@
 <!DOCTYPE html>
-<html>
+<html lang="ES">
 <head>
     <title>Login con Google</title>
 </head>
@@ -8,8 +8,8 @@
 
     <script>
         function iniciarSesion() {
-            var redirectUri = 'http://127.0.0.1:8000/sesioni'; // Reemplaza con la URL a la que se redirigirá después del inicio de sesión
-            var clientId = 'e743f8a36120049fb'; // Reemplaza con tu ID de cliente de Google
+            var redirectUri = 'http://127.0.0.1:8000/oauth'; // Reemplaza con la URL a la que se redirigirá después del inicio de sesión
+            var clientId = '1051565752404-h4un77s3k4el280e589hn98s8ea4bg69.apps.googleusercontent.com'; // Reemplaza con tu ID de cliente de Google
 
             var url = 'https://accounts.google.com/o/oauth2/auth' +
                 '?response_type=code' +