diff --git a/application/api/connector/routes.py b/application/api/connector/routes.py index fb22fe90..5c753f2f 100644 --- a/application/api/connector/routes.py +++ b/application/api/connector/routes.py @@ -1,5 +1,7 @@ +import base64 import datetime import json +import uuid from bson.objectid import ObjectId @@ -13,8 +15,6 @@ from flask import ( from flask_restx import fields, Namespace, Resource - - from application.api.user.tasks import ( ingest_connector_task, ) @@ -246,7 +246,11 @@ class ConnectorAuth(Resource): "status": "pending", "created_at": now }) - state = str(result.inserted_id) + state_dict = { + "provider": provider, + "object_id": str(result.inserted_id) + } + state = base64.urlsafe_b64encode(json.dumps(state_dict).encode()).decode() auth = ConnectorCreator.create_auth(provider) authorization_url = auth.get_authorization_url(state=state) @@ -268,13 +272,15 @@ class ConnectorsCallback(Resource): try: from application.parser.connectors.connector_creator import ConnectorCreator from flask import request, redirect - import uuid - provider = request.args.get('provider', 'google_drive') authorization_code = request.args.get('code') state = request.args.get('state') error = request.args.get('error') + state_dict = json.loads(base64.urlsafe_b64decode(state.encode()).decode()) + provider = state_dict["provider"] + state_object_id = state_dict["object_id"] + if error: if error == "access_denied": return redirect(f"/api/connectors/callback-status?status=cancelled&message=Authentication+was+cancelled.+You+can+try+again+if+you'd+like+to+connect+your+account.&provider={provider}") @@ -285,8 +291,6 @@ class ConnectorsCallback(Resource): if not authorization_code: return redirect(f"/api/connectors/callback-status?status=error&message=Authentication+failed.+Please+try+again+and+make+sure+to+grant+all+requested+permissions.&provider={provider}") - state_object_id = ObjectId(state) - try: auth = ConnectorCreator.create_auth(provider) token_info = auth.exchange_code_for_tokens(authorization_code) @@ -310,7 +314,7 @@ class ConnectorsCallback(Resource): } sessions_collection.find_one_and_update( - {"_id": state_object_id, "provider": provider}, + {"_id": ObjectId(state_object_id), "provider": provider}, { "$set": { "session_token": session_token, diff --git a/application/core/settings.py b/application/core/settings.py index 7ede4e86..1adf8351 100644 --- a/application/core/settings.py +++ b/application/core/settings.py @@ -43,8 +43,7 @@ class Settings(BaseSettings): # Google Drive integration GOOGLE_CLIENT_ID: Optional[str] = None # Replace with your actual Google OAuth client ID GOOGLE_CLIENT_SECRET: Optional[str] = None# Replace with your actual Google OAuth client secret - CONNECTOR_REDIRECT_BASE_URI: Optional[str] = "http://127.0.0.1:7091/api/connectors/callback" - ##append ?provider={provider_name} in your Provider console like http://127.0.0.1:7091/api/connectors/callback?provider=google_drive + CONNECTOR_REDIRECT_BASE_URI: Optional[str] = "http://127.0.0.1:7091/api/connectors/callback" ##add redirect url as it is to your provider's console(gcp) # LLM Cache diff --git a/application/parser/connectors/google_drive/auth.py b/application/parser/connectors/google_drive/auth.py index c282279e..f5fbe056 100644 --- a/application/parser/connectors/google_drive/auth.py +++ b/application/parser/connectors/google_drive/auth.py @@ -23,7 +23,7 @@ class GoogleDriveAuth(BaseConnectorAuth): def __init__(self): self.client_id = settings.GOOGLE_CLIENT_ID self.client_secret = settings.GOOGLE_CLIENT_SECRET - self.redirect_uri = f"{settings.CONNECTOR_REDIRECT_BASE_URI}?provider=google_drive" + self.redirect_uri = f"{settings.CONNECTOR_REDIRECT_BASE_URI}" if not self.client_id or not self.client_secret: raise ValueError("Google OAuth credentials not configured. Please set GOOGLE_CLIENT_ID and GOOGLE_CLIENT_SECRET in settings.")