Update app.py

app.py

@@ -28,8 +28,8 @@ def authenticate_gmail(credentials_file):
             if creds and creds.valid:
                 st.session_state.creds = creds
                 st.session_state.authenticated = True
-                fetch_emails_directly()
-                return
+                st.success("Authentication successful!")
+                return creds
         except Exception as e:
             st.error(f"Invalid token.json file: {e}")
             os.remove('token.json')
@@ -38,8 +38,8 @@ def authenticate_gmail(credentials_file):
         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
+            st.success("Authentication successful!")
+            return st.session_state.creds
         else:
             if not st.session_state.flow:
                 st.session_state.flow = InstalledAppFlow.from_client_secrets_file(credentials_file, SCOPES)
@@ -56,7 +56,6 @@ def authenticate_gmail(credentials_file):
                     st.session_state.flow.fetch_token(code=auth_code)
                     st.session_state.creds = st.session_state.flow.credentials
                     st.session_state.authenticated = True
-                    # 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,
@@ -66,70 +65,37 @@ def authenticate_gmail(credentials_file):
                             "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()
+                    st.success("Authentication successful!")
                 except Exception as e:
                     st.error(f"Error during authentication: {e}")
 
-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.")
+    return st.session_state.creds
 
 # Fetch Emails from Gmail API
 def fetch_emails(service, label):
     emails = []
-    next_page_token = None
-    total_emails = 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
-        ).execute()
-        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')
-            email_data = {
-                "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",
-            }
-            emails.append(email_data)
-
-        progress_bar.progress(min(len(emails) / total_emails, 1.0))
-        next_page_token = results.get('nextPageToken')
-        if not next_page_token:
-            break
-
-    st.success(f"Fetched {len(emails)} emails from {label}")
+    results = service.users().messages().list(userId='me', labelIds=[label], maxResults=100).execute()
+    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')
+        email_data = {
+            "Date": headers.get('Date', ''),
+            "From": headers.get('From', ''),
+            "To": headers.get('To', ''),
+            "Subject": headers.get('Subject', ''),
+            "Snippet": msg.get('snippet', ''),
+            "Body": body,
+        }
+        emails.append(email_data)
+
+    st.success(f"Fetched {len(emails)} emails from {label}.")
     return emails
 
 # Main Page
@@ -153,16 +119,19 @@ def main_page():
         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())
-            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.")
+            creds = authenticate_gmail("credentials.json")
+            if creds:
+                service = build('gmail', 'v1', credentials=creds)
+                label = st.selectbox("Select Label", ["INBOX", "SENT", "DRAFTS", "TRASH", "SPAM"])
+                if st.button("Fetch Emails"):
+                    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
 if "page" not in st.session_state: