from fastapi import FastAPI, Depends, HTTPException, UploadFile, File as FastAPIFile, Request from fastapi.responses import StreamingResponse, RedirectResponse from sqlalchemy.orm import Session from . import crud, models, schemas from .database import SessionLocal, init_db from .services import drive, file_handler from starlette.middleware.sessions import SessionMiddleware from fastapi.middleware.cors import CORSMiddleware from dotenv import load_dotenv import os load_dotenv() app = FastAPI( title="MultiDrive Box", description="Union de varias cuentas de Google Drive para formar un Ășnico pool de almacenamiento.", version="0.1.0" ) # Add CORS middleware to allow cross-origin requests # This is crucial for the frontend (e.g., running on port 3000) # to communicate with the backend (running on port 8000). origins = [ "http://localhost:3000", "http://127.0.0.1:3000", ] app.add_middleware( CORSMiddleware, allow_origins=origins, allow_credentials=True, # Allow cookies to be sent and received allow_methods=["*"], allow_headers=["*"], ) # Add session middleware with a fixed secret key SECRET_KEY = os.getenv("SECRET_KEY") if not SECRET_KEY: raise ValueError("No SECRET_KEY set for session middleware.") app.add_middleware(SessionMiddleware, secret_key=SECRET_KEY) # Dependency def get_db(): db = SessionLocal() try: yield db finally: db.close() @app.on_event("startup") def on_startup(): init_db() @app.get("/api/storage-status", response_model=schemas.StorageStatus) def get_storage_status(db: Session = Depends(get_db)): # Dummy user for now user = crud.get_user(db, 1) if not user: user = crud.create_user(db, schemas.UserCreate(email="dummy@example.com")) accounts = crud.get_accounts(db, user_id=user.id) total_space = sum(acc.drive_space_total for acc in accounts) used_space = sum(acc.drive_space_used for acc in accounts) return { "total_space": total_space, "used_space": used_space, "free_space": total_space - used_space, "accounts": accounts } @app.post("/api/upload-file") async def upload_file(file: UploadFile = FastAPIFile(...), db: Session = Depends(get_db)): # Dummy user user = crud.get_user(db, 1) if not user: user = crud.create_user(db, schemas.UserCreate(email="dummy@example.com")) accounts = crud.get_accounts(db, user_id=user.id) if not accounts: raise HTTPException(status_code=400, detail="No Google Drive accounts linked.") file_path, sha256 = await file_handler.save_temp_file(file) parts = file_handler.split_file(file_path, accounts) uploaded_parts = [] for part in parts: drive_service = drive.get_drive_service(credentials_info=part['account'].credentials) if not drive_service: raise HTTPException(status_code=401, detail=f"Could not get drive service for account {part['account'].google_email}") drive_file_id = drive_service.upload_file(part['path'], part['account']) uploaded_parts.append({ "part_index": part['index'], "size": part['size'], "account_id": part['account'].id, "drive_file_id": drive_file_id, "sha256": part['sha256'] }) download_token = file_handler.generate_token() file_data = schemas.FileCreate( filename=file.filename, original_size=file.size, sha256=sha256, download_token=download_token ) crud.create_file(db, file=file_data, user_id=user.id, parts=uploaded_parts) return {"download_link": f"/api/file/{download_token}"} @app.get("/api/file/{token}", response_model=schemas.File) def get_file_metadata(token: str, db: Session = Depends(get_db)): db_file = crud.get_file_by_token(db, token) if not db_file: raise HTTPException(status_code=404, detail="File not found") return db_file from fastapi.responses import StreamingResponse import shutil import tempfile @app.get("/api/file/{token}/download") def download_file(token: str, db: Session = Depends(get_db)): db_file = crud.get_file_by_token(db, token) if not db_file: raise HTTPException(status_code=404, detail="File not found") temp_dir = tempfile.mkdtemp() part_paths = [] for part in sorted(db_file.parts, key=lambda p: p.part_index): account = crud.get_account(db, part.account_id) drive_service = drive.get_drive_service(credentials_info=account.credentials) if not drive_service: raise HTTPException(status_code=401, detail=f"Could not get drive service for account {account.google_email}") part_path = os.path.join(temp_dir, f"{db_file.filename}.part{part.part_index}") drive.download_file(drive_service, part.drive_file_id, part_path) part_paths.append(part_path) merged_file_path = os.path.join(temp_dir, db_file.filename) file_handler.merge_files(part_paths, merged_file_path) def file_iterator(file_path): with open(file_path, 'rb') as f: yield from f # Clean up temp files shutil.rmtree(temp_dir) return StreamingResponse(file_iterator(merged_file_path), media_type="application/octet-stream", headers={"Content-Disposition": f"attachment; filename={db_file.filename}"}) @app.get("/api/add-account") def add_account(request: Request): print("Received request for /api/add-account") try: redirect_uri = request.url_for('oauth2callback') print(f"Redirect URI for oauth2callback: {redirect_uri}") authorization_url, state = drive.authenticate(redirect_uri) print(f"Generated authorization URL: {authorization_url}") request.session['state'] = state print(f"Stored state in session: {state}") # Directly redirect the user's browser to Google's authorization page return RedirectResponse(authorization_url) except Exception as e: print(f"Error in /api/add-account: {e}") raise HTTPException(status_code=500, detail="Internal server error") @app.get("/api/oauth2callback") def oauth2callback(request: Request, code: str, state: str, db: Session = Depends(get_db)): session_state = request.session.get('state') print(f"Callback received. State from Google: {state}, State from session: {session_state}") if not session_state or state != session_state: print("State mismatch error!") print(f"Session content: {request.session}") raise HTTPException(status_code=400, detail="State mismatch") redirect_uri = request.url_for('oauth2callback') credentials = drive.exchange_code_for_credentials(code, redirect_uri, state) drive_service = drive.get_drive_service(credentials_info=credentials.to_json()) about = drive_service.about().get(fields="user, storageQuota").execute() user_info = about['user'] storage_quota = about['storageQuota'] # Dummy user user = crud.get_user(db, 1) if not user: user = crud.create_user(db, schemas.UserCreate(email="dummy@example.com")) account_data = schemas.AccountCreate( google_email=user_info['emailAddress'], credentials=credentials.to_json(), drive_space_total=int(storage_quota.get('limit', 0)), drive_space_used=int(storage_quota.get('usage', 0)) ) crud.create_account(db, account=account_data, user_id=user.id) return RedirectResponse(url="/")