Update app.py

app.py

@@ -50,7 +50,40 @@ else:
 # Create the MSAL instance for acquiring tokens
 client_instance = get_msal_app()
 
-def get_access_token(code):
+
+def get_access_token(code=None):
+    client_instance = get_msal_app()
+    
+    try:
+        if code:
+            # Initial token acquisition
+            result = client_instance.acquire_token_by_authorization_code(
+                code=code,
+                scopes=SCOPES,
+                redirect_uri=REDIRECT_URI
+            )
+        elif 'refresh_token' in st.session_state:
+            # Token refresh
+            result = client_instance.acquire_token_by_refresh_token(
+                refresh_token=st.session_state['refresh_token'],
+                scopes=SCOPES
+            )
+        else:
+            raise Exception("No authorization code or refresh token available")
+        
+        if 'access_token' in result:
+            st.session_state['access_token'] = result['access_token']
+            if 'refresh_token' in result:
+                st.session_state['refresh_token'] = result['refresh_token']
+            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
+        
+def get_access_token4(code):
     client_instance = get_msal_app()
     
     try:
@@ -128,11 +161,41 @@ def clear_query_params():
     base_url = st.get_page_config().base_url_path
     st.write(f'<meta http-equiv="refresh" content="0; url={base_url}">', unsafe_allow_html=True)
     st.stop()
-    
+
 def process_query_params():
     query_params = st.query_params
     st.write("Debug: All query parameters:", query_params)
 
+    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: {urllib.parse.unquote(error_description)}")
+        st.session_state.clear()
+        st.rerun()
+
+    if 'code' in query_params:
+        code = query_params.get('code')
+        st.write('🔑 Authorization Code Obtained:', code[:10] + '...')
+        
+        try:
+            access_token = get_access_token(code)
+            st.session_state['access_token'] = access_token
+            st.success("Access token acquired successfully!")
+            st.rerun()
+        except Exception as e:
+            if "AADSTS70000" in str(e) and "code has expired" in str(e):
+                st.error("The authorization code has expired. Please log in again.")
+                st.session_state.clear()
+                st.rerun()
+            else:
+                st.error(f"Error acquiring access token: {str(e)}")
+                st.stop()
+    
+def process_query_params3():
+    query_params = st.query_params
+    st.write("Debug: All query parameters:", query_params)
+
     if 'error' in query_params:
         error = query_params.get('error')
         error_description = query_params.get('error_description', 'No description provided')
@@ -232,7 +295,13 @@ def process_query_params2():
         # Display content or image based on the query
         display_content_or_image(query)
 
-
+def get_user_info(access_token):
+    headers = {'Authorization': 'Bearer ' + access_token}
+    response = requests.get('https://graph.microsoft.com/v1.0/me', headers=headers)
+    if response.status_code == 200:
+        return response.json()
+    else:
+        raise Exception(f"Failed to fetch user info: {response.status_code} - {response.text}")
 
 
 # 🏃‍♂️ Main application function
@@ -241,6 +310,31 @@ def main():
 
     process_query_params()
 
+    #added
+    if 'access_token' not in st.session_state:
+        client_instance = get_msal_app()
+        authorization_url = client_instance.get_authorization_request_url(
+            scopes=SCOPES,
+            redirect_uri=REDIRECT_URI
+        )
+        st.write('👋 Please [click here]({}) to log in and authorize the app.'.format(authorization_url))
+        st.stop()
+        
+    try:
+        # Attempt to use the access token
+        user_info = get_user_info(st.session_state['access_token'])
+    except Exception as e:
+        if "token has expired" in str(e).lower():
+            # Token expired, attempt to refresh
+            try:
+                new_token = get_access_token()  # This will use the refresh token if available
+                st.session_state['access_token'] = new_token
+                user_info = get_user_info(new_token)
+            except Exception as refresh_error:
+                st.error("Failed to refresh token. Please log in again.")
+                st.session_state.clear()
+                st.rerun()
+        
     if 'access_token' not in st.session_state:
         query_params = st.query_params
         if 'code' in query_params: