app.py

@@ -98,22 +98,18 @@ def load_data():
 
 # Main function to run the Streamlit app
 def main():
+    st.title('📚 Student Grade Lookup')
+    st.markdown("---")
+    
     if not is_logged_in():
-        st.title('📚 Student Grade Lookup')
-        st.markdown("---")
         st.write("Please log in with your CSPC Google account to access the grade lookup.")
         
         if st.button("Login with Google"):
             flow = create_flow()
-            authorization_url, state = flow.authorization_url(prompt="consent")
-            st.session_state['oauth_state'] = state
-            st.session_state['oauth_state_time'] = time.time()
-            st.experimental_set_query_params(state=state)
+            authorization_url, _ = flow.authorization_url(prompt="consent")
             st.markdown(f"[Login with Google]({authorization_url})")
     else:
-        st.title('📚 Student Grade Lookup')
-        st.markdown("---")
-        
+        st.write("You are logged in.")
         # Load the data
         df = load_data()
 
@@ -184,58 +180,25 @@ def main():
 
         if st.button("Logout"):
             del st.session_state['credentials']
-            st.rerun()
+            st.experimental_rerun()
 
 # Function to handle OAuth callback
 def handle_callback():
     flow = create_flow()
-    query_params = st.experimental_get_query_params()
-    code = query_params.get("code", [None])[0]
-    state = query_params.get("state", [None])[0]
-    
-    if code and state:
-        if 'oauth_state' not in st.session_state or 'oauth_state_time' not in st.session_state:
-            logging.error("OAuth state not found in session")
-            st.error("Authentication failed. Please try again.")
-            return
-
-        if time.time() - st.session_state['oauth_state_time'] > 600:  # 10 minutes expiration
-            logging.error("OAuth state has expired")
-            st.error("Authentication session expired. Please try again.")
-            return
-
-        if state != st.session_state['oauth_state']:
-            logging.error("OAuth state mismatch")
-            st.error("Authentication failed. Please try again.")
-            return
-
-        try:
-            logging.debug(f"Attempting to fetch token with code: {code}")
-            flow.fetch_token(code=code)
-            credentials = flow.credentials
-            st.session_state['credentials'] = credentials.to_json()
-            logging.debug("Token fetch successful")
-            del st.session_state['oauth_state']
-            del st.session_state['oauth_state_time']
-            st.experimental_set_query_params()  # Clear query params
-            return credentials
-        except InvalidGrantError as e:
-            logging.error(f"InvalidGrantError: {str(e)}")
-            st.error("The authorization code has expired or is invalid. Please try logging in again.")
-            if 'credentials' in st.session_state:
-                del st.session_state['credentials']
-        except Exception as e:
-            logging.error(f"Unexpected error during token fetch: {str(e)}")
-            st.error(f"An unexpected error occurred: {str(e)}")
-    else:
-        logging.warning("No authorization code or state found in the URL parameters.")
-        st.error("Authentication failed. Please try again.")
+    try:
+        flow.fetch_token(code=st.experimental_get_query_params()["code"][0])
+        credentials = flow.credentials
+        st.session_state['credentials'] = credentials.to_json()
+        logging.debug("Token fetch successful")
+        st.success("Authentication successful!")
+        return credentials
+    except Exception as e:
+        logging.error(f"Error during authentication: {str(e)}")
+        st.error(f"Authentication failed: {str(e)}")
 
 if __name__ == '__main__':
     logging.debug("Starting the application")
-    query_params = st.experimental_get_query_params()
-    if 'code' in query_params and 'state' in query_params:
-        logging.debug("Authorization code and state found in query parameters")
+    if 'code' in st.experimental_get_query_params():
+        logging.debug("Authorization code found in query parameters")
         handle_callback()
-        st.rerun()
     main()