Update app.py

app.py

@@ -2,140 +2,134 @@ import os
 import streamlit as st
 import requests
 import msal
-import base64
-import hashlib
 import secrets
+import time
+from urllib.parse import urlencode
 
-# 🤓 Load environment variables (Ensure these are set!)
-APPLICATION_ID_KEY = os.getenv('APPLICATION_ID_KEY')
-CLIENT_SECRET_KEY = os.getenv('CLIENT_SECRET_KEY')
-AUTHORITY_URL = 'https://login.microsoftonline.com/common'  # Use 'common' for multi-tenant apps
-REDIRECT_URI = 'https://huggingface.co/spaces/awacke1/MSGraphAPI'  # Update this to match your app's redirect URI
+# Configuration
+APPLICATION_ID = os.getenv('APPLICATION_ID_KEY')
+CLIENT_SECRET = os.getenv('CLIENT_SECRET_KEY')
+AUTHORITY = 'https://login.microsoftonline.com/common'
+REDIRECT_URI = 'https://huggingface.co/spaces/awacke1/MSGraphAPI'
+SCOPES = ['User.Read', 'Calendars.ReadWrite', 'Mail.ReadWrite']
 
-# 🎯 Define the scopes your app will need
-SCOPES = ['User.Read']
+# MSAL setup
+def get_msal_app():
+    return msal.ConfidentialClientApplication(
+        client_id=APPLICATION_ID,
+        client_credential=CLIENT_SECRET,
+        authority=AUTHORITY
+    )
 
-# New function to generate PKCE code verifier and challenge
-def generate_pkce_codes():
-    code_verifier = secrets.token_urlsafe(128)[:128]
-    code_challenge = base64.urlsafe_b64encode(hashlib.sha256(code_verifier.encode()).digest()).decode().rstrip('=')
-    return code_verifier, code_challenge
+# Authentication functions
+def generate_auth_url():
+    msal_app = get_msal_app()
+    state = secrets.token_urlsafe(32)
+    auth_url = msal_app.get_authorization_request_url(
+        scopes=SCOPES,
+        redirect_uri=REDIRECT_URI,
+        state=state
+    )
+    # Store the state in query params
+    new_query_params = st.query_params.to_dict()
+    new_query_params['auth_state'] = state
+    return f"{auth_url}&{urlencode(new_query_params)}"
 
-# 🛠️ Initialize the MSAL client for Public Client App (PKCE support)
-def get_msal_app():
-    return msal.PublicClientApplication(
-        client_id=APPLICATION_ID_KEY,
-        authority=AUTHORITY_URL
+def get_token_from_code(code):
+    msal_app = get_msal_app()
+    result = msal_app.acquire_token_by_authorization_code(
+        code=code,
+        scopes=SCOPES,
+        redirect_uri=REDIRECT_URI
     )
+    if 'access_token' in result:
+        return result
+    else:
+        raise Exception(f"Error acquiring token: {result.get('error_description')}")
 
-# 🔐 Acquire access token using authorization code
-def get_access_token(code, code_verifier=None):
-    client_instance = get_msal_app()
-    
-    st.write("Debug: MSAL App Configuration:")
-    st.write(f"Client ID: {APPLICATION_ID_KEY[:5]}...")
-    st.write(f"Authority: {AUTHORITY_URL}")
-    st.write(f"Redirect URI: {REDIRECT_URI}")
-    
-    try:
-        # Attempt to acquire token, use PKCE code_verifier if provided
-        result = client_instance.acquire_token_by_authorization_code(
-            code=code,
-            scopes=SCOPES,
-            redirect_uri=REDIRECT_URI,
-            code_verifier=code_verifier  # Include only if PKCE is enabled
-        )
-        
-        if 'access_token' in result:
-            return result['access_token']
-        else:
-            error_description = result.get('error_description', 'No error description provided')
-            raise Exception(f"Error acquiring token: {error_description}")
-    except Exception as e:
-        st.error(f"Exception in get_access_token: {str(e)}")
-        raise
+# API call function
+def make_api_call(endpoint, token):
+    headers = {'Authorization': f'Bearer {token}'}
+    response = requests.get(f'https://graph.microsoft.com/v1.0/{endpoint}', headers=headers)
+    if response.status_code == 200:
+        return response.json()
+    else:
+        st.error(f"API call failed: {response.status_code} - {response.text}")
+        return None
 
-# 🏃‍♂️ Main function to process the query parameters and handle the token exchange
-def process_query_params():
-    try:
-        query_params = st.experimental_get_query_params()
-        st.write("Debug: All query parameters:", query_params)
+# Main application
+def main():
+    st.title("🦄 MS Graph API Integration")
 
-        if 'error' in query_params:
-            error = query_params.get('error')
-            error_description = query_params.get('error_description', 'No description provided')
-            st.error(f"Authentication Error: {error}")
-            st.error(f"Error Description: {error_description}")
-            st.stop()
+    # Debug information
+    st.sidebar.write("Debug Info:")
+    st.sidebar.write(f"Query Params: {st.query_params.to_dict()}")
 
-        if 'code' in query_params:
-            code = query_params.get('code')[0]  # MS Graph returns the code as a list
-            st.write('🔑 Authorization Code Obtained:', code[:10] + '...')
-            
-            try:
-                # Retrieve code_verifier from session state
-                code_verifier = st.session_state.get('code_verifier')
-                if not code_verifier:
-                    st.error("Code verifier not found in session state.")
-                    st.stop()
-                
-                # Acquire the access token
-                access_token = get_access_token(code, code_verifier)
-                st.session_state['access_token'] = access_token
-                st.success("Access token acquired successfully!")
-                
-                # Clear the query parameters from the URL
-                st.experimental_set_query_params()
-                st.experimental_rerun()
-            except Exception as e:
-                st.error(f"Error acquiring access token: {str(e)}")
-                st.stop()
-        else:
-            st.warning("No authorization code found in the query parameters.")
-    except Exception as e:
-        st.error(f"Error processing query parameters: {str(e)}")
-        st.stop()
+    if 'code' in st.query_params and 'state' in st.query_params:
+        received_state = st.query_params['state']
+        expected_state = st.query_params.get('auth_state')
+        
+        if received_state != expected_state:
+            st.error(f"Invalid state parameter. Expected {expected_state}, got {received_state}")
+            st.error("Please try logging in again.")
+            st.query_params.clear()
+            st.rerun()
+        
+        try:
+            token = get_token_from_code(st.query_params['code'])
+            st.session_state['token'] = token
+            st.query_params.clear()
+            st.success("Successfully authenticated!")
+            st.rerun()
+        except Exception as e:
+            st.error(f"Authentication failed: {str(e)}")
+            st.query_params.clear()
+            st.rerun()
 
-# Main application function
-def main():
-    st.title("🦄 MS Graph API with AI & Cloud Integration with M365")
+    if 'token' not in st.session_state:
+        auth_url = generate_auth_url()
+        st.write("Please log in to continue:")
+        st.markdown(f"[Login with Microsoft]({auth_url})")
+        return
 
-    if 'access_token' not in st.session_state:
-        if 'code_verifier' not in st.session_state:
-            # Generate PKCE codes and store code_verifier in session
-            code_verifier, code_challenge = generate_pkce_codes()
-            st.session_state['code_verifier'] = code_verifier
-        else:
-            code_verifier = st.session_state['code_verifier']
-            code_challenge = base64.urlsafe_b64encode(hashlib.sha256(code_verifier.encode()).digest()).decode().rstrip('=')
+    # User is authenticated, show the main app
+    token = st.session_state['token']
+    st.sidebar.success("Authenticated successfully!")
+    
+    # Display user info
+    user_info = make_api_call('me', token['access_token'])
+    if user_info:
+        st.sidebar.write(f"Welcome, {user_info.get('displayName', 'User')}!")
 
-        # Get MSAL app and construct the authorization URL
-        client_instance = get_msal_app()
-        auth_url = client_instance.get_authorization_request_url(
-            scopes=SCOPES,
-            redirect_uri=REDIRECT_URI,
-            code_challenge=code_challenge,
-            code_challenge_method="S256"
-        )
-        st.write('👋 Please [click here]({}) to log in and authorize the app.'.format(auth_url))
-        st.stop()
+    # App functionality
+    option = st.sidebar.selectbox(
+        "Choose a function",
+        ["View Emails", "View Calendar", "View OneDrive Files"]
+    )
 
-    # Process query parameters to acquire token
-    process_query_params()
+    if option == "View Emails":
+        emails = make_api_call('me/messages?$top=10', token['access_token'])
+        if emails:
+            for email in emails['value']:
+                st.write(f"Subject: {email['subject']}")
+                st.write(f"From: {email['from']['emailAddress']['name']}")
+                st.write("---")
 
-    # If access token is present, greet the user
-    if 'access_token' in st.session_state:
-        access_token = st.session_state['access_token']
-        headers = {'Authorization': 'Bearer ' + access_token}
+    elif option == "View Calendar":
+        events = make_api_call('me/events?$top=10', token['access_token'])
+        if events:
+            for event in events['value']:
+                st.write(f"Event: {event['subject']}")
+                st.write(f"Start: {event['start']['dateTime']}")
+                st.write("---")
 
-        response = requests.get('https://graph.microsoft.com/v1.0/me', headers=headers)
-        if response.status_code == 200:
-            user_info = response.json()
-            st.write(f"👋 Hello, {user_info['displayName']}!")
-        else:
-            st.error("Failed to fetch user info.")
-            st.write(response.text)
+    elif option == "View OneDrive Files":
+        files = make_api_call('me/drive/root/children', token['access_token'])
+        if files:
+            for file in files['value']:
+                st.write(f"File: {file['name']}")
+                st.write(f"Type: {file['file']['mimeType'] if 'file' in file else 'Folder'}")
+                st.write("---")
 
-# 🚀 Run the main function
 if __name__ == "__main__":
-    main()
+    main()