Update app.py

app.py

@@ -2,134 +2,140 @@ import os
 import streamlit as st
 import requests
 import msal
+import base64
+import hashlib
 import secrets
-import time
-from urllib.parse import urlencode
 
-# 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']
+# 🤓 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
 
-# MSAL setup
-def get_msal_app():
-    return msal.ConfidentialClientApplication(
-        client_id=APPLICATION_ID,
-        client_credential=CLIENT_SECRET,
-        authority=AUTHORITY
-    )
+# 🎯 Define the scopes your app will need
+SCOPES = ['User.Read']
 
-# 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)}"
+# 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
 
-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
+# 🛠️ 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
     )
-    if 'access_token' in result:
-        return result
-    else:
-        raise Exception(f"Error acquiring token: {result.get('error_description')}")
 
-# 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
+# 🔐 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
 
-# Main application
-def main():
-    st.title("🦄 MS Graph API Integration")
+# 🏃‍♂️ 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)
 
-    # Debug information
-    st.sidebar.write("Debug Info:")
-    st.sidebar.write(f"Query Params: {st.query_params.to_dict()}")
+        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()
 
-    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()
+        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 '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
+# Main application function
+def main():
+    st.title("🦄 MS Graph API with AI & Cloud Integration with M365")
 
-    # 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')}!")
+    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('=')
 
-    # App functionality
-    option = st.sidebar.selectbox(
-        "Choose a function",
-        ["View Emails", "View Calendar", "View OneDrive Files"]
-    )
+        # 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()
 
-    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("---")
+    # Process query parameters to acquire token
+    process_query_params()
 
-    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("---")
+    # 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 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("---")
+        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)
 
+# 🚀 Run the main function
 if __name__ == "__main__":
-    main()
+    main()