Update app.py

app.py

@@ -164,7 +164,6 @@ def clear_query_params():
 
 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')
@@ -174,23 +173,25 @@ def process_query_params():
         st.session_state.clear()
         st.rerun()
 
-    if 'code' in query_params:
+    if 'code' in query_params and 'state' in query_params:
+        received_state = query_params.get('state')
+        if received_state != st.session_state.get('auth_state'):
+            st.error("Invalid state parameter. Please try logging in again.")
+            st.session_state.clear()
+            st.rerun()
+
         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!")
+            token = get_access_token(code)
+            st.session_state['access_token'] = token
+            st.success("Successfully authenticated!")
+            # Clear the URL parameters
+            st.experimental_set_query_params()
             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()
+            st.error(f"Error acquiring access token: {str(e)}")
+            st.session_state.clear()
+            st.rerun()
     
 def process_query_params3():
     query_params = st.query_params
@@ -303,6 +304,30 @@ def get_user_info(access_token):
     else:
         raise Exception(f"Failed to fetch user info: {response.status_code} - {response.text}")
 
+def initiate_auth_flow():
+    client_instance = get_msal_app()
+    auth_url = client_instance.get_authorization_request_url(
+        scopes=SCOPES,
+        redirect_uri=REDIRECT_URI,
+        state=generate_state()
+    )
+    st.write('👋 Please [click here]({}) to log in and authorize the app.'.format(auth_url))
+
+def generate_state():
+    state = secrets.token_urlsafe(32)
+    st.session_state['auth_state'] = state
+    return state
+
+def is_token_valid(token):
+    if not token:
+        return False
+    try:
+        # Make a simple API call to check if the token is still valid
+        headers = {'Authorization': f'Bearer {token}'}
+        response = requests.get('https://graph.microsoft.com/v1.0/me', headers=headers)
+        return response.status_code == 200
+    except:
+        return False
 
 # 🏃‍♂️ Main application function
 def main():
@@ -310,6 +335,17 @@ def main():
 
     process_query_params()
 
+    
+    # Check if we have a valid access token
+    if 'access_token' not in st.session_state or not is_token_valid(st.session_state.get('access_token')):
+        # If not, initiate the login process
+        initiate_auth_flow()
+        st.stop()
+
+    # If we have a valid token, proceed with the app
+    access_token = st.session_state['access_token']
+
+    
     #added
     if 'access_token' not in st.session_state:
         client_instance = get_msal_app()