import uuid import aiohttp from fastapi import FastAPI, HTTPException, Request, UploadFile from fastapi.responses import FileResponse from fastapi import Depends, FastAPI from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer from pyrate_limiter import AbstractBucket, BucketFactory, Duration, InMemoryBucket, Limiter, MonotonicClock, Rate, RateItem from fastapi_limiter.depends import RateLimiter from fastapi import Depends import asyncmy import asyncio import os from dotenv import load_dotenv from contextlib import asynccontextmanager import logging load_dotenv() log_level = os.getenv("LOG_LEVEL", "INFO").upper() log_formatter = logging.Formatter( fmt="%(asctime)s [%(levelname)s] %(message)s", datefmt="%Y-%m-%d %H:%M:%S", ) file_handler = logging.FileHandler("data/api.log", encoding="utf-8") file_handler.setFormatter(log_formatter) console_handler = logging.StreamHandler() console_handler.setFormatter(log_formatter) logging.basicConfig( level=getattr(logging, log_level, logging.INFO), handlers=[file_handler, console_handler], ) security = HTTPBearer() SECRET = os.getenv("SECRET") @asynccontextmanager async def connect_db(app: FastAPI): app.state.pool = await asyncmy.create_pool( host=os.getenv("DB_HOST", "localhost"), port=int(os.getenv("DB_PORT", 3306)), user=os.getenv("DB_USER"), password=os.getenv("DB_PASSWORD"), db=os.getenv("DB_NAME"), minsize=5, maxsize=20 ) await create_tables(app.state.pool) task = asyncio.create_task(fetch_images()) try: yield finally: task.cancel() try: await task except asyncio.CancelledError: pass app.state.pool.close() await app.state.pool.wait_closed() class MultiBucketFactory(BucketFactory): def __init__(self, rates, clock): self.clock = clock self.rates = rates self.buckets = {} def wrap_item(self, name: str, weight: int = 1) -> RateItem: """Time-stamping item, return a RateItem""" now = self.clock.now() return RateItem(name, now, weight=weight) def get(self, item: RateItem) -> AbstractBucket: if item.name not in self.buckets: new_bucket = self.create(InMemoryBucket, self.rates) self.buckets.update({item.name: new_bucket}) return self.buckets[item.name] app = FastAPI(lifespan=connect_db) rates = [Rate(50, Duration.MINUTE)] limiter = Limiter(MultiBucketFactory(rates,MonotonicClock())) async def create_tables(pool): async with pool.acquire() as conn: async with conn.cursor() as cursor: await cursor.execute(""" CREATE TABLE IF NOT EXISTS images ( id INT AUTO_INCREMENT PRIMARY KEY, url VARCHAR(255) NOT NULL, filename VARCHAR(255) NOT NULL, source VARCHAR(255) NOT NULL, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) """) await conn.commit() @app.get("/") async def root(): return {"message": "yes the api works, maybe i will create a small landing page later here"} @app.get("/random", dependencies=[Depends(RateLimiter(limiter=limiter))]) async def get_random_bnuy(request: Request): async with app.state.pool.acquire() as conn: async with conn.cursor() as cursor: await cursor.execute("SELECT filename, source, url FROM images ORDER BY RAND() LIMIT 1;") result = await cursor.fetchone() if result: filepath = os.path.join("data/images", result[0]) if os.path.exists(filepath): return {"url": f"{request.base_url}images/{result[0]}", "source": result[1], "original_url": result[2]} else: raise HTTPException(status_code=404, detail="Image file not found") else: raise HTTPException(status_code=404, detail="No available images found") @app.get("/images/{filename}", dependencies=[Depends(RateLimiter(limiter=limiter))]) async def get_image(filename: str): async with app.state.pool.acquire() as conn: async with conn.cursor() as cursor: await cursor.execute("SELECT filename FROM images WHERE filename = %s", (filename,)) result = await cursor.fetchone() if result: filepath = os.path.join("data/images", result[0]) if os.path.exists(filepath): return FileResponse(filepath) else: raise HTTPException(status_code=404, detail="Image file not found") else: raise HTTPException(status_code=404, detail="Image not found") async def fetch_images(): from collector import save_picture while True: try: logging.info("Starting image collection...") await save_picture(app.state.pool) logging.info("Image collection completed. Sleeping for 1 hour...") except Exception as e: logging.error(f"Error during image collection: {e}") await asyncio.sleep(86400) # Sleep for 24 hours @app.post("/upload", dependencies=[Depends(RateLimiter(limiter=limiter))]) async def upload_image(file: UploadFile, request: Request, credentials: HTTPAuthorizationCredentials = Depends(security)): # if SECRET isn't set, return Unauthrorized if credentials.credentials != SECRET or credentials.credentials is None: raise HTTPException(status_code=401, detail="Unauthorized") if not file: raise HTTPException(status_code=400, detail="No file uploaded") if not file.filename.lower().endswith(('.jpg', '.jpeg', '.png', '.gif')) or file.content_type not in ["image/jpeg", "image/png", "image/gif"]: raise HTTPException(status_code=400, detail="Unsupported file type") try: content = await file.read() generate_filename = str(uuid.uuid4()) + os.path.splitext(file.filename)[1] filename = os.path.join("data/images", generate_filename) with open(filename, "wb") as f: f.write(content) logging.info(f"Saved uploaded image to {filename}") async with app.state.pool.acquire() as conn: async with conn.cursor() as cursor: await cursor.execute( "INSERT INTO images (url, filename, source) VALUES (%s, %s, %s)", (f"{request.base_url}images/{generate_filename}", generate_filename, "user_upload") ) await conn.commit() except Exception as e: logging.error(f"Error saving uploaded image: {e}") raise HTTPException(status_code=500, detail="Failed to save image")