Implement OAuth 2.0 authentication flow and session management; remove token.pickle file

app/controllers/mail.py

@@ -6,7 +6,7 @@ from datetime import datetime, timedelta
 from venv import logger
 from ics import Calendar
 
-import pandas as pd
+# import pandas as pd
 from langchain_core.documents import Document
 from langchain_community.document_loaders import (
     PyPDFLoader,
@@ -16,7 +16,7 @@ from langchain_community.document_loaders import (
 )
 
 from models.db import vectorstore
-from models.mails import build_gmail_service
+# from models.mails import build_gmail_service
 
 SCOPES = ["https://www.googleapis.com/auth/gmail.readonly"]
 EMAIL_PATTERN = r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}"
@@ -24,10 +24,8 @@ EMAIL_PATTERN = r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}"
 ATTACHMENTS_DIR = "cache"
 os.makedirs(ATTACHMENTS_DIR, exist_ok=True)
 
-service = build_gmail_service()
-
-
-def search_emails(query):
+# service = build_gmail_service()
+def search_emails(service, query):
     """Search emails based on a query."""
     result = service.users().messages().list(userId="me", q=query).execute()
     messages = []
@@ -43,13 +41,13 @@ def search_emails(query):
     return messages
 
 
-def list_emails(messages):
+def list_emails(service, messages):
     """
     Processes a list of email messages, extracts metadata, decodes content, and handles attachments.
 
     Args:
         messages (list): A list of email message dictionaries, where each dictionary contains
-                         at least an 'id' key representing the email's unique identifier.
+                        at least an 'id' key representing the email's unique identifier.
 
     Returns:
         None: The function processes the emails and adds the extracted documents to a vector store.
@@ -94,7 +92,7 @@ def list_emails(messages):
         )
         metadata["user_id"] = service.users().getProfile(userId="me").execute().get("emailAddress")
         metadata["msg_id"] = msg["id"]
-        print(metadata, msg["payload"]["mimeType"])
+        # print(metadata, msg["payload"]["mimeType"])
         ids = []
         documents = []
         mime_types = []
@@ -199,7 +197,7 @@ def list_emails(messages):
             vectorstore.add_documents(documents=documents, ids=ids)
 
 
-def collect(query=(datetime.today() - timedelta(days=21)).strftime("after:%Y/%m/%d")):
+def collect(service, query=(datetime.today() - timedelta(days=21)).strftime("after:%Y/%m/%d")):
     """
     Main function to search and list emails from Gmail.
 
@@ -211,47 +209,47 @@ def collect(query=(datetime.today() - timedelta(days=21)).strftime("after:%Y/%m/
         None
     """
     query = "subject:Re: Smartcareers algorithm debug and improvement'"
-    emails = search_emails(query)
+    emails = search_emails(service, query)
     if emails:
         print("Found %d emails:\n", len(emails))
         logger.info("Found %d emails after two_weeks_ago:\n", len(emails))
-        list_emails(emails)
+        list_emails(service, emails)
         logger.info("Listing emails...")
         return f"{len(emails)} emails added to the collection."
     else:
         logger.info("No emails found after two weeks ago.")
 
 
-def get_documents():
-    """
-    Main function to list emails from the database.
+# def get_documents(self):
+#     """
+#     Main function to list emails from the database.
 
-    This function lists all emails stored in the database.
+#     This function lists all emails stored in the database.
 
-    Returns:
-        None
-    """
-    data = vectorstore.get()
-    df = pd.DataFrame(
-        {"ids": data["ids"], "documents": data["documents"], "metadatas": data["metadatas"]}
-    )
-    df.to_excel("collection_data.xlsx", index=False)
-    df = pd.concat(
-        [df.drop("metadatas", axis=1), df["metadatas"].apply(pd.Series)], axis=1
-    ).to_excel("collection_data_expand.xlsx", index=False)
+#     Returns:
+#         None
+#     """
+#     data = vectorstore.get()
+#     df = pd.DataFrame(
+#         {"ids": data["ids"], "documents": data["documents"], "metadatas": data["metadatas"]}
+#     )
+#     df.to_excel("collection_data.xlsx", index=False)
+#     df = pd.concat(
+#         [df.drop("metadatas", axis=1), df["metadatas"].apply(pd.Series)], axis=1
+#     ).to_excel("collection_data_expand.xlsx", index=False)
 
 
-def get():
-    """
-    Main function to list emails from the database.
+# def get(self):
+#     """
+#     Main function to list emails from the database.
 
-    This function lists all emails stored in the database.
+#     This function lists all emails stored in the database.
 
-    Returns:
-        None
-    """
-    data = vectorstore.get()
-    df = pd.DataFrame(
-        {"id": data["ids"], "documents": data["documents"], "metadatas": data["metadatas"]}
-    )
-    return df.to_dict(orient="records")
+#     Returns:
+#         None
+#     """
+#     data = vectorstore.get()
+#     df = pd.DataFrame(
+#         {"id": data["ids"], "documents": data["documents"], "metadatas": data["metadatas"]}
+#     )
+#     return df.to_dict(orient="records")

app/main.py

@@ -1,14 +1,72 @@
 """Module to handle the main FastAPI application and its endpoints."""
 import logging
-from fastapi import FastAPI
+
+from fastapi import FastAPI, Request
 from fastapi.middleware.cors import CORSMiddleware
-from router import content, mail
+from jose import jwt
+from router import auth, content, mail
+from starlette.middleware.base import BaseHTTPMiddleware
+
+SECRET_KEY = "your-secret-key"
+ALGORITHM = "HS256"
+
+class SessionMiddleware(BaseHTTPMiddleware):
+    """
+    Middleware to manage session data using JWT (JSON Web Tokens) stored in cookies.
+
+    This middleware intercepts incoming HTTP requests to extract session data from a 
+    "session_token" cookie. If the cookie exists and contains a valid JWT, the session 
+    data is decoded and attached to the request's state. If the cookie is missing or 
+    invalid, an empty session is initialized.
 
+    After processing the request, the middleware encodes the updated session data into 
+    a new JWT and sets it as a "session_token" cookie in the response.
+
+    Attributes:
+        SECRET_KEY (str): The secret key used to encode and decode the JWT.
+        ALGORITHM (str): The algorithm used for encoding and decoding the JWT.
+
+    Methods:
+        dispatch(request: Request, call_next): Intercepts the request to manage session 
+        data and modifies the response to include the updated session token.
+
+    Cookie Parameters:
+        session_token (str): A JWT containing session data. This cookie is HTTP-only 
+        and has a maximum age of 3600 seconds (1 hour).
+
+    Raises:
+        jwt.JWTError: If the session token cannot be decoded due to invalid signature 
+        or other issues.
+    """
+    async def dispatch(self, request: Request, call_next):
+        session_token = request.cookies.get("session_token")
+        if session_token:
+            try:
+                session_data = jwt.decode(session_token, SECRET_KEY, algorithms=[ALGORITHM])
+            except jwt.JWTError:
+                session_data = {}
+        else:
+            session_data = {}
+        request.state.session = session_data
+        response = await call_next(request)
+        session_token = jwt.encode(request.state.session, SECRET_KEY, algorithm=ALGORITHM)
+        response.set_cookie(
+            key="session_token",
+            value=session_token,
+            httponly=True,
+            max_age=3600
+        )
+        return response
+
+logging.basicConfig(
+    format='%(asctime)s - %(levelname)s - %(funcName)s - %(message)s')
+logging.getLogger().setLevel(logging.ERROR)
 
 app = FastAPI(docs_url="/")
 
 app.include_router(content.router)
 app.include_router(mail.router)
+app.include_router(auth.router)
 
 origins = [
     "*"
@@ -21,11 +79,7 @@ app.add_middleware(
     allow_methods=["*"],
     allow_headers=["*"],
 )
-
-logging.basicConfig(
-    format='%(asctime)s - %(levelname)s - %(funcName)s - %(message)s')
-logging.getLogger().setLevel(logging.ERROR)
-
+app.add_middleware(SessionMiddleware)
 
 @app.get("/_health")
 def health():

app/router/auth.py

@@ -0,0 +1,75 @@
+"""Module for defining the main routes of the API."""
+import os
+import json
+import pickle
+from fastapi import APIRouter, Request
+from fastapi.responses import JSONResponse
+from google_auth_oauthlib.flow import InstalledAppFlow
+from googleapiclient.discovery import build
+
+router = APIRouter(tags=["auth"])
+
+CLIENT_ID = os.environ.get("CLIENT_ID")
+CLIENT_SECRET = os.environ.get("CLIENT_SECRET")
+REDIRECT_URI = os.environ.get("REDIRECT_URI")
+
+SCOPES = ["https://www.googleapis.com/auth/gmail.readonly"]
+
+# Client config for OAuth flow
+CLIENT_CONFIG = {
+    "web": {
+        "client_id": CLIENT_ID,
+        "client_secret": CLIENT_SECRET,
+        "redirect_uris": [REDIRECT_URI],
+        "auth_uri": "https://accounts.google.com/o/oauth2/auth",
+        "token_uri": "https://oauth2.googleapis.com/token",
+    }
+}
+
+@router.get("/auth/google/url")
+async def get_auth_url():
+    """
+    Handles the generation of a Google OAuth 2.0 authorization URL.
+
+    This endpoint initializes an OAuth 2.0 flow using the provided client configuration
+    and scopes, sets the redirect URI, and generates an authorization URL for the user
+    to grant access.
+
+    Returns:
+        dict: A dictionary containing the generated authorization URL under the key "url".
+    """
+    flow = InstalledAppFlow.from_client_config(CLIENT_CONFIG, SCOPES)
+    flow.redirect_uri = REDIRECT_URI
+    auth_url, _ = flow.authorization_url(access_type="offline", prompt="consent")
+    return JSONResponse({"url": auth_url})
+
+@router.get("/auth/google/callback")
+async def google_callback(code: str, request: Request):
+    flow = InstalledAppFlow.from_client_config(CLIENT_CONFIG, SCOPES)
+    flow.redirect_uri = REDIRECT_URI
+    flow.fetch_token(code=code)
+    credentials = flow.credentials
+    request.state.session["credential"] = json.loads(credentials.to_json())
+    # cred_dict = (request.state.session.get("credential"))
+    # cred = Credentials(
+    #     token=cred_dict["token"],
+    #     refresh_token=cred_dict["refresh_token"],
+    #     token_uri=cred_dict["token_uri"],
+    #     client_id=cred_dict["client_id"],
+    #     client_secret=cred_dict["client_secret"],
+    #     scopes=cred_dict["scopes"],
+    # )
+    # service = build("gmail", "v1", credentials=Credentials(
+    #     token=cred_dict["token"],
+    #     refresh_token=cred_dict["refresh_token"],
+    #     token_uri=cred_dict["token_uri"],
+    #     client_id=cred_dict["client_id"],
+    #     client_secret=cred_dict["client_secret"],
+    #     scopes=cred_dict["scopes"],
+    # ))
+    service = build("gmail", "v1", credentials=credentials)
+    profile = service.users().getProfile(userId="me").execute()
+    print(({"profile": profile}))
+    with open(f"{profile['emailAddress']}.pickle", "wb") as token:
+        pickle.dump(credentials, token)
+    return JSONResponse(profile)

app/router/mail.py

@@ -1,12 +1,17 @@
 """Module for defining the main routes of the API."""
-from fastapi import APIRouter
+import os
+import pickle
+from fastapi import APIRouter, Request
 from fastapi.responses import JSONResponse
+
 from controllers import mail
+from google.oauth2.credentials import Credentials
+from googleapiclient.discovery import build
 
 router = APIRouter(prefix="/mail", tags=["mail"])
 
 @router.post("")
-def collect():
+def collect(emailAddress: str, request: Request):
     """
     Handles the chat POST request.
 
@@ -16,19 +21,34 @@ def collect():
     Returns:
         str: The generated response from the chat function.
     """
-    mail.collect()
+    # emailAddress = req['emailAddress']
+    if os.path.exists(f"{emailAddress}.pickle"):
+        with open(f"{emailAddress}.pickle", "rb") as token:
+            credentials = pickle.load(token)
+    else:
+        cred_dict = request.state.session.get("credential")
+        credentials = Credentials(
+            token=cred_dict["token"],
+            refresh_token=cred_dict["refresh_token"],
+            token_uri=cred_dict["token_uri"],
+            client_id=cred_dict["client_id"],
+            client_secret=cred_dict["client_secret"],
+            scopes=cred_dict["scopes"],
+        )
+    mailservice = build("gmail", "v1", credentials=credentials)
+    mail.collect(mailservice)
     return JSONResponse(content={"message": "Mail collected successfully."})
 
-@router.get("")
-def get():
-    """
-    Handles the chat POST request.
+# @router.get("")
+# def get():
+#     """
+#     Handles the chat POST request.
 
-    Args:
-        query (ReqData): The request data containing the query parameters.
+#     Args:
+#         query (ReqData): The request data containing the query parameters.
 
-    Returns:
-        str: The generated response from the chat function.
-    """
-    result = mail.get()
-    return JSONResponse(content= result)
+#     Returns:
+#         str: The generated response from the chat function.
+#     """
+#     # result = mail.get()
+#     return JSONResponse(content= result)