File size: 5,713 Bytes
fc4bf5d 7a66f6d 8d50900 7a66f6d fc4bf5d 7a66f6d 4dc60dc fc4bf5d 7a66f6d 3a85b14 8d50900 fc4bf5d 03b52dc fc4bf5d 8d50900 7a66f6d 8d50900 7a66f6d 8d50900 fc4bf5d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 |
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 """
<html>
<head>
<title>Proxy Usage</title>
</head>
<body>
<h1>Proxy Usage</h1>
<p>To use the proxy, make a GET request to the following URL:</p>
<p><code>/{api_key}?url={your_url}</code></p>
<p>Replace <code>{api_key}</code> with your API key and <code>{your_url}</code> with the URL you want to proxy.</p>
<h2>Example with cURL:</h2>
<p><code>curl -X GET "http://your-hugging-face-space-url/your_api_key?url=https://example.com"</code></p>
</body>
</html>
"""
@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))
|