Fix authentication callback issues: add cache control headers, prevent double callbacks, improve error handling, and ensure server-side sessions work properly

app.py

@@ -1,4 +1,4 @@
-from flask import Flask, request, render_template, send_file, flash, redirect, url_for, jsonify, session
+from flask import Flask, request, render_template, send_file, flash, redirect, url_for, jsonify, session, make_response
 import os
 import tempfile
 import shutil
@@ -27,7 +27,17 @@ app.config['MAX_CONTENT_LENGTH'] = 1000 * 1024 * 1024 * 1024  # 1000TB
 app.config['SESSION_TYPE'] = 'filesystem'
 app.config['SESSION_FILE_DIR'] = '/tmp/flask_sessions'
 app.config['SESSION_PERMANENT'] = False
-Session(app)
+
+# Ensure session directory exists and is writable
+os.makedirs('/tmp/flask_sessions', exist_ok=True)
+
+try:
+    Session(app)
+    print("DEBUG: Flask-Session initialized successfully")
+except Exception as e:
+    print(f"ERROR: Failed to initialize Flask-Session: {e}")
+    # Fallback to client-side sessions if server-side fails
+    pass
 
 # Load environment variables
 load_dotenv()
@@ -136,8 +146,22 @@ def auth_callback():
     print(f"DEBUG: Callback URL: {request.url}")
     print(f"DEBUG: Session keys: {list(session.keys())}")
     
+    # Add cache control headers to prevent browser caching
+    response = make_response()
+    response.headers['Cache-Control'] = 'no-cache, no-store, must-revalidate'
+    response.headers['Pragma'] = 'no-cache'
+    response.headers['Expires'] = '0'
+    
+    # Check if already authenticated to prevent double callback
+    if 'oauth_token' in session:
+        print("DEBUG: Already authenticated, redirecting to index")
+        response.set_data(redirect(url_for('index')).get_data())
+        return response
+    
     if 'oauth_state' not in session:
-        return redirect(url_for('login'))
+        print("DEBUG: No oauth_state in session, redirecting to login")
+        response.set_data(redirect(url_for('login')).get_data())
+        return response
     
     # Parse the authorization response
     from urllib.parse import parse_qs, urlparse
@@ -148,38 +172,31 @@ def auth_callback():
     
     if 'error' in query_params:
         print(f"DEBUG: Error from Azure AD: {query_params['error'][0]}")
-        return f"Authentication error: {query_params['error'][0]}", 400
+        response.set_data(f"Authentication error: {query_params['error'][0]}")
+        response.status_code = 400
+        return response
     
     if 'code' not in query_params or 'state' not in query_params:
-        return "Invalid authorization response", 400
+        print("DEBUG: Invalid authorization response - missing code or state")
+        response.set_data("Invalid authorization response")
+        response.status_code = 400
+        return response
     
     code = query_params['code'][0]
     state = query_params['state'][0]
     
     if state != session['oauth_state']:
-        return "State mismatch", 400
+        print(f"DEBUG: State mismatch: {state} != {session['oauth_state']}")
-    
+        response.set_data("State mismatch")
-    # Use OAuth2Session for token exchange to maintain PKCE consistency
+        response.status_code = 400
-    oauth_session = OAuth2Session(
+        return response
-        client_id=AZURE_CLIENT_ID,
-        scope=AZURE_SCOPES,
-        redirect_uri=AZURE_REDIRECT_URI
-    )
-    
-    # Set the same PKCE that was used for authorization
-    if 'code_verifier' in session and hasattr(oauth_session, '_client'):
-        oauth_session._client.code_verifier = session['code_verifier']
-    
-    # Set the authorization response for the session
-    redirect_uri_with_params = f"{AZURE_REDIRECT_URI}?code={code}&state={state}"
-    print(f"DEBUG: Redirect URI with params: {redirect_uri_with_params}")
     
     # Exchange code for token using direct HTTP request with PKCE
     import requests
     
     token_data = {
         'client_id': AZURE_CLIENT_ID,
-        'client_secret': AZURE_CLIENT_SECRET,  # Include client_secret for confidential client
+        'client_secret': AZURE_CLIENT_SECRET,
         'code': code,
         'grant_type': 'authorization_code',
         'redirect_uri': AZURE_REDIRECT_URI,
@@ -193,26 +210,46 @@ def auth_callback():
     
     print(f"DEBUG: Token request data keys: {list(token_data.keys())}")
     
-    token_response = requests.post(TOKEN_URL, data=token_data)
+    try:
-    print(f"DEBUG: Token response status: {token_response.status_code}")
+        token_response = requests.post(TOKEN_URL, data=token_data, timeout=30)
-    print(f"DEBUG: Token response: {token_response.text[:200]}...")
+        print(f"DEBUG: Token response status: {token_response.status_code}")
+        print(f"DEBUG: Token response: {token_response.text[:200]}...")
         
-    if token_response.status_code != 200:
+        if token_response.status_code != 200:
-        return f"Token exchange failed: {token_response.text}", 400
+            print(f"DEBUG: Token exchange failed: {token_response.text}")
+            response.set_data(f"Token exchange failed: {token_response.text}")
+            response.status_code = 400
+            return response
         
-    token = token_response.json()
+        token = token_response.json()
-    print(f"DEBUG: Token received successfully")
+        print(f"DEBUG: Token received successfully")
-    session['oauth_token'] = token
+        session['oauth_token'] = token
         
-    # Get user info
+        # Get user info
-    user_info = oauth_session.get('https://graph.microsoft.com/v1.0/me').json()
+        headers = {'Authorization': f"Bearer {token['access_token']}"}
-    session['user_info'] = user_info
+        user_response = requests.get('https://graph.microsoft.com/v1.0/me', headers=headers, timeout=30)
         
-    # Clean up
+        if user_response.status_code == 200:
-    session.pop('oauth_state', None)
+            user_info = user_response.json()
-    session.pop('code_verifier', None)
+            session['user_info'] = user_info
+            print(f"DEBUG: User info retrieved: {user_info.get('displayName', 'Unknown')}")
+        else:
+            print(f"DEBUG: Failed to get user info: {user_response.text}")
+            session['user_info'] = {'displayName': 'Unknown User'}
         
-    return redirect(url_for('index'))
+        # Clean up session
+        session.pop('oauth_state', None)
+        session.pop('code_verifier', None)
+        
+        print("DEBUG: Authentication successful, redirecting to index")
+        response.set_data(redirect(url_for('index')).get_data())
+        return response
+        
+    except requests.exceptions.RequestException as e:
+        print(f"DEBUG: Request error during token exchange: {e}")
+        response.set_data(f"Authentication service error: {str(e)}")
+        response.status_code = 500
+        return response
 
 @app.route('/logout')
 def logout():