app.py

@@ -11,10 +11,15 @@ from google.auth.transport.requests import Request
 import json
 import logging
 from oauthlib.oauth2.rfc6749.errors import InvalidGrantError
+import secrets
+import time
 
 # Load environment variables from .env file
 load_dotenv()
 
+# Set up logging
+logging.basicConfig(level=logging.DEBUG)
+
 # Set page configuration
 st.set_page_config(page_title="Student Grade Lookup", page_icon="📚", layout="centered")
 
@@ -100,7 +105,9 @@ def main():
         
         if st.button("Login with Google"):
             flow = create_flow()
-            authorization_url, _ = flow.authorization_url(prompt="consent")
+            authorization_url, state = flow.authorization_url(prompt="consent")
+            st.session_state['oauth_state'] = state
+            st.session_state['oauth_state_time'] = time.time()
             st.markdown(f"[Login with Google]({authorization_url})")
     else:
         st.title('📚 Student Grade Lookup')
@@ -182,13 +189,32 @@ def main():
 def handle_callback():
     flow = create_flow()
     code = st.query_params.get("code")
-    if code:
+    state = st.query_params.get("state")
+    
+    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']
             return credentials
         except InvalidGrantError as e:
             logging.error(f"InvalidGrantError: {str(e)}")
@@ -199,16 +225,13 @@ def handle_callback():
             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 found in the URL parameters.")
-        st.error("No authorization code found in the URL parameters.")
-
-# Set up logging
-logging.basicConfig(level=logging.DEBUG)
+        logging.warning("No authorization code or state found in the URL parameters.")
+        st.error("Authentication failed. Please try again.")
 
 if __name__ == '__main__':
     logging.debug("Starting the application")
-    if 'code' in st.query_params:
-        logging.debug("Authorization code found in query parameters")
+    if 'code' in st.query_params and 'state' in st.query_params:
+        logging.debug("Authorization code and state found in query parameters")
         handle_callback()
         st.rerun()
     main()