from fastapi import FastAPI, HTTPException, BackgroundTasks, Request
from fastapi.responses import HTMLResponse, StreamingResponse
from curl_cffi import requests
import os
import time
import hashlib
import threading
from dotenv import load_dotenv
from collections import defaultdict
load_dotenv()
app = FastAPI()
# --- Rate Limiting Configuration ---
RATE_LIMIT_REQUESTS = 20
RATE_LIMIT_WINDOW = 60 # 60 seconds
rate_limit_data = defaultdict(list)
# --- Cache Configuration ---
CACHE_DIR = "cache"
CACHE_MAX_SIZE = 100 * 1024 * 1024 # 100 MB
CACHE_MAX_AGE = 3600 # 1 hour in seconds
MAX_FILE_SIZE = 10 * 1024 * 1024 # 10 MB
CACHE_CLEANUP_INTERVAL = 600 # 10 minutes
# --- Cache Initialization ---
if not os.path.exists(CACHE_DIR):
os.makedirs(CACHE_DIR)
# --- Cache Cleanup Functions ---
def clean_cache():
"""Cleans the cache directory based on size and age.
This function is designed to be run in a background thread.
"""
while True:
try:
total_size = 0
files = []
for filename in os.listdir(CACHE_DIR):
filepath = os.path.join(CACHE_DIR, filename)
if not os.path.isfile(filepath):
continue
stat = os.stat(filepath)
total_size += stat.st_size
files.append((filepath, stat.st_mtime, stat.st_size))
# 1. Remove expired files
now = time.time()
for filepath, mtime, _ in files:
if now - mtime > CACHE_MAX_AGE:
os.remove(filepath)
print(f"Removed expired cache file: {filepath}")
# 2. Enforce total cache size limit (LRU)
if total_size > CACHE_MAX_SIZE:
# Sort files by last modification time (oldest first)
files.sort(key=lambda x: x[1])
while total_size > CACHE_MAX_SIZE and files:
filepath, _, size = files.pop(0)
os.remove(filepath)
total_size -= size
print(f"Removed old cache file to free space: {filepath}")
except Exception as e:
print(f"Error during cache cleanup: {e}")
time.sleep(CACHE_CLEANUP_INTERVAL)
# --- Background Thread for Cache Cleanup ---
cache_cleanup_thread = threading.Thread(target=clean_cache, daemon=True)
cache_cleanup_thread.start()
# --- API Endpoints ---
@app.get("/", response_class=HTMLResponse)
async def root():
return """
Proxy Usage
Proxy Usage
To use the proxy, make a GET request to the following URL:
/{api_key}?url={your_url}
Replace {api_key} with your API key and {your_url} with the URL you want to proxy.
Example with cURL:
curl -X GET "http://your-hugging-face-space-url/your_api_key?url=https://example.com"
"""
@app.get("/{api_key}")
async def proxy(api_key: str, url: str, request: Request, background_tasks: BackgroundTasks):
# --- Rate Limiting Logic ---
client_ip = request.headers.get("X-Forwarded-For")
if client_ip:
client_ip = client_ip.split(',')[0].strip()
else:
client_ip = request.headers.get("X-Real-IP")
if not client_ip:
client_ip = request.client.host
now = time.time()
# Remove old timestamps
rate_limit_data[client_ip] = [t for t in rate_limit_data[client_ip] if now - t < RATE_LIMIT_WINDOW]
# Check request limit
if len(rate_limit_data[client_ip]) >= RATE_LIMIT_REQUESTS:
raise HTTPException(status_code=429, detail="Too Many Requests")
# Add current timestamp
rate_limit_data[client_ip].append(now)
# --- API Key Check ---
expected_api_key = os.environ.get("API_KEY")
if not expected_api_key or api_key != expected_api_key:
raise HTTPException(status_code=401, detail="Invalid API key")
url_hash = hashlib.md5(url.encode()).hexdigest()
cache_path = os.path.join(CACHE_DIR, url_hash)
# Check for valid cached file
if os.path.exists(cache_path):
if now - os.path.getmtime(cache_path) < CACHE_MAX_AGE:
def file_iterator(file_path, chunk_size=8192):
with open(file_path, "rb") as f:
while True:
chunk = f.read(chunk_size)
if not chunk:
break
yield chunk
return StreamingResponse(file_iterator(cache_path))
# If not cached or expired, fetch from URL
try:
response = requests.get(url, impersonate="chrome110", stream=True, timeout=15)
response.raise_for_status() # Raise an exception for bad status codes
async def stream_and_cache():
total_size = 0
do_cache = True
temp_cache_path = cache_path + ".tmp"
with open(temp_cache_path, "wb") as f:
for chunk in response.iter_content(chunk_size=8192):
total_size += len(chunk)
if total_size > MAX_FILE_SIZE:
do_cache = False
if do_cache:
f.write(chunk)
yield chunk
if do_cache:
os.rename(temp_cache_path, cache_path)
else:
os.remove(temp_cache_path)
return StreamingResponse(stream_and_cache())
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))