Update app.py

app.py

@@ -10,146 +10,163 @@ import base64
 
 SCOPES = ['https://www.googleapis.com/auth/gmail.readonly']
 
-# Helper Functions
-def save_credentials_to_file(creds):
-    """Save credentials to token.json file."""
-    with open("token.json", "w") as token_file:
-        json.dump({
-            "token": creds.token,
-            "refresh_token": creds.refresh_token,
-            "token_uri": creds.token_uri,
-            "client_id": creds.client_id,
-            "client_secret": creds.client_secret,
-            "scopes": creds.scopes,
-        }, token_file)
-
+# Initialize session state variables
+if "authenticated" not in st.session_state:
+    st.session_state.authenticated = False
+if "creds" not in st.session_state:
+    st.session_state.creds = None
+if "auth_url" not in st.session_state:
+    st.session_state.auth_url = None
+if "flow" not in st.session_state:
+    st.session_state.flow = None
+
+# Authenticate Gmail API
 def authenticate_gmail(credentials_file):
-    """Authenticate with Gmail API."""
-    creds = None
-    if os.path.exists("token.json"):
+    if os.path.exists('token.json'):
         try:
-            creds = Credentials.from_authorized_user_file("token.json", SCOPES)
+            creds = Credentials.from_authorized_user_file('token.json', SCOPES)
+            if creds and creds.valid:
+                st.session_state.creds = creds
+                st.session_state.authenticated = True
+                fetch_emails_directly()
+                return
         except Exception as e:
             st.error(f"Invalid token.json file: {e}")
-            os.remove("token.json")
-
-    if not creds or not creds.valid:
-        if creds and creds.expired and creds.refresh_token:
-            creds.refresh(Request())
+            os.remove('token.json')
+
+    if not st.session_state.creds or not st.session_state.creds.valid:
+        if st.session_state.creds and st.session_state.creds.expired and st.session_state.creds.refresh_token:
+            st.session_state.creds.refresh(Request())
+            st.session_state.authenticated = True
+            fetch_emails_directly()
+            return
         else:
-            flow = InstalledAppFlow.from_client_secrets_file(credentials_file, SCOPES)
-            flow.redirect_uri = "http://localhost"
-            auth_url, _ = flow.authorization_url(prompt="consent")
+            if not st.session_state.flow:
+                st.session_state.flow = InstalledAppFlow.from_client_secrets_file(credentials_file, SCOPES)
+                st.session_state.flow.redirect_uri = 'http://localhost'
+                auth_url, _ = st.session_state.flow.authorization_url(prompt='consent')
+                st.session_state.auth_url = auth_url
+
             st.info("Please visit this URL to authorize the application:")
-            st.code(auth_url)
+            st.code(st.session_state.auth_url)
 
             auth_code = st.text_input("Enter the authorization code:")
-            if st.button("Submit Authorization Code"):
+            if st.button("Submit Authentication Code"):
                 try:
-                    flow.fetch_token(code=auth_code)
-                    creds = flow.credentials
-                    save_credentials_to_file(creds)
-                    st.success("Authentication successful!")
+                    st.session_state.flow.fetch_token(code=auth_code)
+                    st.session_state.creds = st.session_state.flow.credentials
                     st.session_state.authenticated = True
-                    st.session_state.creds = creds
-                    st.experimental_rerun()
+                    # Save the credentials to a file for later use
+                    with open('token.json', 'w') as token_file:
+                        json.dump({
+                            "token": st.session_state.creds.token,
+                            "refresh_token": st.session_state.creds.refresh_token,
+                            "token_uri": st.session_state.creds.token_uri,
+                            "client_id": st.session_state.creds.client_id,
+                            "client_secret": st.session_state.creds.client_secret,
+                            "scopes": st.session_state.creds.scopes
+                        }, token_file)
+                    st.success("Authentication successful! Fetching emails...")
+                    fetch_emails_directly()
                 except Exception as e:
                     st.error(f"Error during authentication: {e}")
 
-    return creds
-
+def fetch_emails_directly():
+    if st.session_state.authenticated and st.session_state.creds:
+        service = build('gmail', 'v1', credentials=st.session_state.creds)
+        label = "INBOX"  # Default label, you can make this dynamic
+        st.info(f"Fetching emails from {label}...")
+        emails = fetch_emails(service, label)
+        if emails:
+            df = pd.DataFrame(emails)
+            st.dataframe(df)
+            csv = df.to_csv(index=False).encode('utf-8')
+            st.download_button("Download Emails as CSV", csv, f"{label}_emails.csv", "text/csv")
+        else:
+            st.warning(f"No emails found in {label}.")
+    else:
+        st.error("You are not authenticated. Please authenticate again.")
 
+# Fetch Emails from Gmail API
 def fetch_emails(service, label):
-    """Fetch emails from Gmail API."""
     emails = []
     next_page_token = None
     total_emails = 0
 
-    # Get total email count in the selected label
-    results = service.users().messages().list(userId="me", labelIds=[label], maxResults=1).execute()
-    total_emails = results.get("resultSizeEstimate", 0)
+    # Get total emails
+    results = service.users().messages().list(userId='me', labelIds=[label], maxResults=1).execute()
+    total_emails = results.get('resultSizeEstimate', 0)
     st.info(f"Total emails in {label}: {total_emails}")
 
     progress_bar = st.progress(0)
 
     while True:
         results = service.users().messages().list(
-            userId="me", labelIds=[label], maxResults=100, pageToken=next_page_token
+            userId='me', labelIds=[label], maxResults=100, pageToken=next_page_token
         ).execute()
-        messages = results.get("messages", [])
+        messages = results.get('messages', [])
         for message in messages:
-            msg = service.users().messages().get(userId="me", id=message["id"], format="full").execute()
-            headers = {header["name"]: header["value"] for header in msg["payload"]["headers"]}
-            body = ""
-            if "parts" in msg["payload"]:
-                for part in msg["payload"]["parts"]:
-                    if part["mimeType"] == "text/plain":
-                        body = base64.urlsafe_b64decode(part["body"].get("data", "").encode("UTF-8")).decode("UTF-8")
+            msg = service.users().messages().get(userId='me', id=message['id'], format='full').execute()
+            headers = {header['name']: header['value'] for header in msg['payload']['headers']}
+            body = ''
+            if 'parts' in msg['payload']:
+                for part in msg['payload']['parts']:
+                    if part['mimeType'] == 'text/plain':
+                        body = base64.urlsafe_b64decode(part['body'].get('data', '').encode('UTF-8')).decode('UTF-8')
             email_data = {
-                "Date": headers.get("Date", ""),
-                "From": headers.get("From", ""),
-                "To": headers.get("To", ""),
-                "Subject": headers.get("Subject", ""),
-                "Snippet": msg.get("snippet", ""),
+                "Date": headers.get('Date', ''),
+                "From": headers.get('From', ''),
+                "To": headers.get('To', ''),
+                "Subject": headers.get('Subject', ''),
+                "Snippet": msg.get('snippet', ''),
                 "Body": body,
-                "Attachments": "Yes" if "attachmentId" in str(msg["payload"]) else "No",
+                "Attachments": "Yes" if 'attachmentId' in str(msg['payload']) else "No",
             }
             emails.append(email_data)
 
         progress_bar.progress(min(len(emails) / total_emails, 1.0))
-        next_page_token = results.get("nextPageToken")
+        next_page_token = results.get('nextPageToken')
         if not next_page_token:
             break
 
     st.success(f"Fetched {len(emails)} emails from {label}")
     return emails
 
-
-# Pages
+# Main Page
 def main_page():
-    """Main Page for Authentication."""
-    st.title("Gmail Email Fetcher - Authenticate")
-    st.sidebar.header("Navigation")
-    st.sidebar.radio("Go to:", ["Main Page"], key="page_selector")
-
+    st.sidebar.title("Steps to Use")
+    st.sidebar.write("""
+    1. Enable Gmail API from [Google Cloud Console](https://console.cloud.google.com/).
+    2. Download the `credentials.json` file.
+    3. Enter your Gmail email and App Password.
+    4. Upload your `credentials.json` file.
+    5. Authenticate and proceed to fetch emails.
+    """)
+
+    st.title("Gmail Email Fetcher")
+
+    user_email = st.text_input("Gmail Email")
+    app_password = st.text_input("App Password", type="password")
     credentials_file = st.file_uploader("Upload credentials.json", type="json")
 
-    if st.button("Start Authentication"):
-        if not credentials_file:
-            st.error("Please upload the credentials.json file.")
+    if st.button("Authenticate"):
+        if not user_email or not app_password or not credentials_file:
+            st.error("Please provide all required inputs (Gmail, App Password, credentials.json)")
         else:
             # Save credentials.json locally
             with open("credentials.json", "wb") as f:
                 f.write(credentials_file.getbuffer())
-            creds = authenticate_gmail("credentials.json")
-            if creds:
-                st.session_state.authenticated = True
-                st.session_state.creds = creds
-                st.experimental_rerun()
+            authenticate_gmail("credentials.json")
 
+    # Display Authentication Status
+    if st.session_state.authenticated:
+        st.success("You are authenticated!")
+    else:
+        st.warning("You are not authenticated yet.")
 
-def email_fetch_page():
-    """Page to Fetch Emails."""
-    st.title("Gmail Email Fetcher - Fetch Emails")
-    st.sidebar.header("Navigation")
-    st.sidebar.radio("Go to:", ["Email Fetch Page"], key="page_selector")
+# Navigation
+if "page" not in st.session_state:
+    st.session_state.page = "main"
 
-    if not st.session_state.authenticated:
-        st.error("You are not authenticated. Please go back to the Main Page.")
-    else:
-        label = st.selectbox("Select Label", ["INBOX", "SENT", "DRAFTS", "TRASH", "SPAM"])
-        if st.button("Fetch Emails"):
-            service = build("gmail", "v1", credentials=st.session_state.creds)
-            emails = fetch_emails(service, label)
-            if emails:
-                df = pd.DataFrame(emails)
-                st.dataframe(df)
-                csv = df.to_csv(index=False).encode("utf-8")
-                st.download_button("Download Emails as CSV", csv, f"{label}_emails.csv", "text/csv")
-
-
-# Navigation Logic
-if "authenticated" in st.session_state and st.session_state.authenticated:
-    email_fetch_page()
-else:
+if st.session_state.page == "main":
     main_page()