Update main.py

main.py

@@ -14,6 +14,7 @@ import pandas as pd
 import requests
 from urllib.parse import urlparse, unquote
 import datetime
+from datetime import timedelta
 
 
 app = Flask(__name__)
@@ -256,17 +257,27 @@ def get_spending_overview():
 def verify_admin(auth_header):
     if not auth_header or not auth_header.startswith('Bearer '):
         raise ValueError('Invalid token')
-    
+
     token = auth_header.split(' ')[1]
     uid = verify_token(token)
     if not uid:
         raise PermissionError('Invalid user')
-    
+
     user_ref = db.reference(f'users/{uid}')
     user_data = user_ref.get()
     if not user_data or not user_data.get('is_admin', False):
         raise PermissionError('Admin access required')
 
+    # **Set custom claim here AFTER admin verification is successful:**
+    try:
+        auth.set_custom_user_claims(uid, {"admin": True})
+        print(f"Custom admin claim set for user {uid}") # Optional log
+    except Exception as e:
+        print(f"Error setting custom admin claim: {e}") # Log any errors
+        raise PermissionError('Error setting admin claim, but admin verified') # Or handle error as needed
+
+    return uid # Return uid as before
+
 # ========================================
 # New Admin Endpoints
 # ========================================
@@ -428,10 +439,17 @@ def view_receipt(transaction_id):
         if not blob_path:
             return jsonify({'error': 'Could not determine blob path from URL'}), 500
 
+        print(f"Blob path for view: {blob_path}")  # Debug log
+
         blob = bucket.blob(blob_path)
-        # Generate a signed URL valid for 10 minutes
-        signed_url = blob.generate_signed_url(expiration=datetime.timedelta(minutes=10))
+        if not blob.exists():
+            print("Blob does not exist at path:", blob_path)
+            return jsonify({'error': 'Blob not found'}), 404
+
+        signed_url = blob.generate_signed_url(expiration=timedelta(minutes=10))
+        print(f"Signed URL for view: {signed_url}")  # Debug log
         r = requests.get(signed_url)
+        print(f"View endpoint response status: {r.status_code}")  # Debug log
         if r.status_code != 200:
             return jsonify({'error': 'Unable to fetch image from storage'}), 500
 
@@ -457,13 +475,20 @@ def download_receipt(transaction_id):
         if not blob_path:
             return jsonify({'error': 'Could not determine blob path from URL'}), 500
 
+        print(f"Blob path for download: {blob_path}")  # Debug log
+
         blob = bucket.blob(blob_path)
-        signed_url = blob.generate_signed_url(expiration=datetime.timedelta(minutes=10))
+        if not blob.exists():
+            print("Blob does not exist at path:", blob_path)
+            return jsonify({'error': 'Blob not found'}), 404
+
+        signed_url = blob.generate_signed_url(expiration=timedelta(minutes=10))
+        print(f"Signed URL for download: {signed_url}")  # Debug log
         r = requests.get(signed_url)
+        print(f"Download endpoint response status: {r.status_code}")  # Debug log
         if r.status_code != 200:
             return jsonify({'error': 'Unable to fetch image from storage'}), 500
 
-        # Use attachment_filename if you're on Flask 1.x; if using Flask 2.x, you can replace it with download_name.
         return send_file(
             io.BytesIO(r.content),
             mimetype='image/jpeg',
@@ -509,7 +534,6 @@ def delete_user(uid):
     except Exception as e:
         return jsonify({'error': str(e)}), 500
 
-# ... (rest of the code remains the same)
 
 if __name__ == '__main__':
     app.run(debug=True, host="0.0.0.0", port=7860)