from fastapi import FastAPI, HTTPException, Header, Request from pyppeteer import launch import os import asyncio import json import re import sqlite3 import time from datetime import datetime, timedelta from typing import Optional, Dict, List, Any from urllib.parse import unquote from apscheduler.schedulers.background import BackgroundScheduler from apscheduler.triggers.cron import CronTrigger from fastapi.middleware.base import BaseHTTPMiddleware # Add imports for browser pool from asyncio import Queue, Lock from contextlib import asynccontextmanager app = FastAPI() # Get API key from environment variable API_KEY = os.getenv('API_KEY') if not API_KEY: raise ValueError("API_KEY environment variable must be set") # Get cache expiry time from environment variable (default: 36 hours) CACHE_EXPIRY_HOURS = int(os.getenv('CACHE_EXPIRY_HOURS', '36')) # Get cleanup cron schedule from environment variable (default: every day at 3 AM) CLEANUP_CRON = os.getenv('CLEANUP_CRON', '0 3 * * *') # Define custom user agent CUSTOM_USER_AGENT = 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/69.0.3497.100 Safari/537.36' # Browser pool configuration MAX_BROWSERS = int(os.getenv('MAX_BROWSERS', '5')) # Maximum number of browser instances BROWSER_TTL = int(os.getenv('BROWSER_TTL', '3600')) # Time to live for browser instances in seconds browser_pool = Queue() browser_lock = Lock() browser_creation_times = {} # Rate limiting configuration RATE_LIMIT_MINUTE = int(os.getenv('RATE_LIMIT_MINUTE', '60')) # requests per minute RATE_LIMIT_WINDOW = 60 # window size in seconds class RateLimitMiddleware(BaseHTTPMiddleware): def __init__(self, app): super().__init__(app) self.requests = {} self.lock = asyncio.Lock() async def dispatch(self, request: Request, call_next): # Skip rate limiting for health check if request.url.path == "/" and request.method == "HEAD": return await call_next(request) api_key = request.headers.get("x-api-key") if not api_key: raise HTTPException(status_code=401, detail="API key required") async with self.lock: now = time.time() # Clean old requests self.requests = {k: v for k, v in self.requests.items() if now - v[-1] < RATE_LIMIT_WINDOW} # Get request times for this API key requests = self.requests.get(api_key, []) # Remove old requests outside the window requests = [t for t in requests if now - t < RATE_LIMIT_WINDOW] if len(requests) >= RATE_LIMIT_MINUTE: oldest = requests[0] wait_time = RATE_LIMIT_WINDOW - (now - oldest) raise HTTPException( status_code=429, detail=f"Rate limit exceeded. Try again in {int(wait_time)} seconds" ) requests.append(now) self.requests[api_key] = requests return await call_next(request) # Add rate limiting middleware app.add_middleware(RateLimitMiddleware) async def create_browser(): """Create a new browser instance""" browser = await launch( headless=True, executablePath='/usr/bin/google-chrome', args=[ '--no-sandbox', '--disable-setuid-sandbox', '--disable-dev-shm-usage', '--disable-accelerated-2d-canvas', '--disable-gpu', '--disable-extensions', '--disable-sync', '--disable-background-networking', '--disable-default-apps', '--disable-translate', '--disable-background-timer-throttling', '--disable-backgrounding-occluded-windows', '--disable-client-side-phishing-detection', '--disable-features=site-per-process', '--disable-hang-monitor', '--disable-ipc-flooding-protection', '--disable-popup-blocking', '--disable-prompt-on-repost', '--disable-renderer-backgrounding', '--memory-pressure-off', '--no-first-run', '--safebrowsing-disable-auto-update', ], handleSIGINT=False, handleSIGTERM=False, handleSIGHUP=False, ignoreHTTPSErrors=True ) browser_creation_times[browser] = time.time() return browser async def check_browser_health(): """Check browser health and recycle if needed""" while True: try: # Sleep for 5 minutes between checks await asyncio.sleep(300) async with browser_lock: # Get all browsers from the pool browsers = [] while not browser_pool.empty(): browsers.append(await browser_pool.get()) # Check each browser for browser in browsers: try: # Check if browser is too old if time.time() - browser_creation_times.get(browser, 0) > BROWSER_TTL: await browser.close() del browser_creation_times[browser] browser = await create_browser() else: # Quick health check await browser.pages() # Put back in pool if healthy await browser_pool.put(browser) except Exception: # If unhealthy, close and create new try: await browser.close() except: pass if browser in browser_creation_times: del browser_creation_times[browser] new_browser = await create_browser() await browser_pool.put(new_browser) except Exception as e: print(f"Error in browser health check: {str(e)}") # Initialize browser pool @app.on_event("startup") async def init_browser_pool(): """Initialize the browser pool with some browsers""" for _ in range(min(3, MAX_BROWSERS)): # Start with 3 browsers or MAX_BROWSERS, whichever is smaller browser = await create_browser() await browser_pool.put(browser) print(f"Browser pool initialized with {browser_pool.qsize()} browsers") # Start browser health check task asyncio.create_task(check_browser_health()) @app.on_event("shutdown") async def cleanup_browser_pool(): """Clean up all browsers in the pool""" while not browser_pool.empty(): try: browser = await browser_pool.get_nowait() await browser.close() if browser in browser_creation_times: del browser_creation_times[browser] except: pass # Initialize SQLite database def init_db(): global DB_PATH # Try to use the mounted volume first db_path = '/db/cache.db' db_dir = os.path.dirname(db_path) # Check if directory exists and is writable dir_writable = False if os.path.exists(db_dir): try: test_file = os.path.join(db_dir, '.write_test') with open(test_file, 'w') as f: f.write('test') os.remove(test_file) dir_writable = True except (IOError, PermissionError): print(f"Directory {db_dir} exists but is not writable") dir_writable = False # If directory doesn't exist or isn't writable, try to create it if not os.path.exists(db_dir) or not dir_writable: try: os.makedirs(db_dir, exist_ok=True) # Test if we can write to the directory test_file = os.path.join(db_dir, '.write_test') with open(test_file, 'w') as f: f.write('test') os.remove(test_file) print(f"Created directory: {db_dir}") dir_writable = True except Exception as e: print(f"Warning: Could not create or write to directory {db_dir}: {e}") # Fallback to using a local database file db_path = 'cache.db' print(f"Using local database file: {db_path}") try: conn = sqlite3.connect(db_path) cursor = conn.cursor() cursor.execute(''' CREATE TABLE IF NOT EXISTS cache ( url TEXT, route TEXT, data TEXT, timestamp INTEGER, PRIMARY KEY (url, route) ) ''') conn.commit() conn.close() print(f"Database initialized at {db_path}") # Update the global DB_PATH DB_PATH = db_path except sqlite3.OperationalError as e: print(f"Error initializing database at {db_path}: {e}") # Fallback to using a local database file if the mounted volume has permission issues db_path = 'cache.db' print(f"Falling back to local database file: {db_path}") try: conn = sqlite3.connect(db_path) cursor = conn.cursor() cursor.execute(''' CREATE TABLE IF NOT EXISTS cache ( url TEXT, route TEXT, data TEXT, timestamp INTEGER, PRIMARY KEY (url, route) ) ''') conn.commit() conn.close() print(f"Local database initialized at {db_path}") # Update the global DB_PATH DB_PATH = db_path except sqlite3.OperationalError as e2: print(f"Error initializing local database: {e2}") raise # Define the database path DB_PATH = '/db/cache.db' # Get cached data if it exists and is not older than the expiry time def get_cached_data(url, route): conn = sqlite3.connect(DB_PATH) cursor = conn.cursor() cache_expiry = int(time.time()) - (CACHE_EXPIRY_HOURS * 60 * 60) # Convert hours to seconds cursor.execute( "SELECT data FROM cache WHERE url = ? AND route = ? AND timestamp > ?", (url, route, cache_expiry) ) result = cursor.fetchone() conn.close() if result: print(f"Cache hit for {url} on route {route}") return json.loads(result[0]) return None # Save data to cache def save_to_cache(url, route, data): conn = sqlite3.connect(DB_PATH) cursor = conn.cursor() timestamp = int(time.time()) # Convert data to JSON string data_json = json.dumps(data) cursor.execute( "INSERT OR REPLACE INTO cache (url, route, data, timestamp) VALUES (?, ?, ?, ?)", (url, route, data_json, timestamp) ) conn.commit() conn.close() print(f"Saved to cache: {url} on route {route}") # Function to clean up old cache entries def cleanup_old_cache_entries(): try: print(f"Running scheduled cache cleanup (entries older than {CACHE_EXPIRY_HOURS} hours, pagination: {CACHE_EXPIRY_HOURS * 31} hours)") conn = sqlite3.connect(DB_PATH) cursor = conn.cursor() # Calculate the timestamp for entries older than the expiry time expiry_timestamp = int(time.time()) - (CACHE_EXPIRY_HOURS * 60 * 60) pagination_expiry_timestamp = int(time.time()) - (CACHE_EXPIRY_HOURS * 31 * 60 * 60) # Get count of entries to be deleted (non-pagination) cursor.execute("SELECT COUNT(*) FROM cache WHERE route != 'pagination' AND timestamp < ?", (expiry_timestamp,)) count_non_pagination = cursor.fetchone()[0] # Get count of pagination entries to be deleted cursor.execute("SELECT COUNT(*) FROM cache WHERE route = 'pagination' AND timestamp < ?", (pagination_expiry_timestamp,)) count_pagination = cursor.fetchone()[0] # Delete old non-pagination entries cursor.execute("DELETE FROM cache WHERE route != 'pagination' AND timestamp < ?", (expiry_timestamp,)) # Delete old pagination entries cursor.execute("DELETE FROM cache WHERE route = 'pagination' AND timestamp < ?", (pagination_expiry_timestamp,)) conn.commit() conn.close() print(f"Cache cleanup completed: {count_non_pagination} non-pagination entries and {count_pagination} pagination entries removed") except Exception as e: print(f"Error during cache cleanup: {e}") # Initialize scheduler for periodic cache cleanup scheduler = BackgroundScheduler() scheduler.add_job( cleanup_old_cache_entries, CronTrigger.from_crontab(CLEANUP_CRON), id='cache_cleanup_job', replace_existing=True ) # Initialize database on startup init_db() # Start the scheduler when the application starts @app.on_event("startup") def start_scheduler(): scheduler.start() print(f"Cache cleanup scheduler started with cron: {CLEANUP_CRON}") print(f"Cache entries will expire after {CACHE_EXPIRY_HOURS} hours") # Shutdown the scheduler when the application stops @app.on_event("shutdown") def shutdown_scheduler(): scheduler.shutdown(wait=False) print("Cache cleanup scheduler stopped") async def wait_for_network_idle(page): """Wait until no network requests are in flight""" await page.waitForNetworkIdle(idleTime=500, timeout=30000) @app.head("/") async def health_check(): return {"status": "ok"} async def safe_browser_operation(url, operation_func): """Safely perform a browser operation with proper cleanup""" async with get_browser() as browser: try: # Create a new page page = await browser.newPage() # Set reasonable viewport await page.setViewport({'width': 1280, 'height': 800}) # Set user agent await page.setUserAgent(CUSTOM_USER_AGENT) # Set reasonable timeout page.setDefaultNavigationTimeout(30000) # Enable request interception to block unnecessary resources await page.setRequestInterception(True) async def intercept(request): # Block unnecessary resource types if request.resourceType in ['image', 'media', 'font', 'stylesheet']: await request.abort() else: await request.continue_() page.on('request', lambda req: asyncio.ensure_future(intercept(req))) # Perform the operation result = await operation_func(page) return result except Exception as e: print(f"Error during browser operation: {str(e)}") raise finally: try: # Ensure page is properly closed if 'page' in locals(): await page.close() except Exception as e: print(f"Error closing page: {str(e)}") @app.get("/") async def visit_url(url: str, x_api_key: Optional[str] = Header(None)): # Validate API key if not x_api_key or x_api_key != API_KEY: raise HTTPException(status_code=401, detail="Invalid API key") # Decode URL if it's encoded decoded_url = unquote(url) # Check cache first cached_result = get_cached_data(decoded_url, "visit") if cached_result: return cached_result try: print(f"Visiting URL: {decoded_url}") # Define the operation to perform with the browser async def visit_operation(page): try: response = await page.goto(decoded_url, waitUntil='networkidle2', timeout=30000) if not response: print(f"Warning: No response object returned for {decoded_url}") # Get page content content = await page.content() return {"status": "success", "content": content} except Exception as e: print(f"Error during page navigation: {e}") # Try to get content anyway try: content = await page.content() return {"status": "partial", "content": content, "error": str(e)} except: raise HTTPException(status_code=500, detail=f"Failed to get page content: {str(e)}") # Perform the operation result = await safe_browser_operation(decoded_url, visit_operation) # Save to cache save_to_cache(decoded_url, "visit", result) return result except Exception as e: print(f"Error visiting URL {decoded_url}: {e}") raise HTTPException(status_code=500, detail=str(e)) @app.get("/seo") async def extract_seo(url: str, x_api_key: Optional[str] = Header(None)): """Extract SEO information from a website""" # Validate API key if not x_api_key or x_api_key != API_KEY: raise HTTPException(status_code=401, detail="Invalid API key") # Decode URL if it's encoded decoded_url = unquote(url) # Check cache first cached_result = get_cached_data(decoded_url, "seo") if cached_result: return cached_result try: print(f"Extracting SEO from: {decoded_url}") # Define the operation to perform with the browser async def seo_operation(page): try: response = await page.goto(decoded_url, waitUntil='networkidle2', timeout=30000) # Extract SEO information seo_data = await page.evaluate('''() => { const data = { title: document.title || '', description: '', canonical: '', h1: [], h2: [], images: 0, links: 0 }; // Get meta description const metaDescription = document.querySelector('meta[name="description"]'); if (metaDescription) { data.description = metaDescription.getAttribute('content') || ''; } // Get canonical link const canonicalLink = document.querySelector('link[rel="canonical"]'); if (canonicalLink) { data.canonical = canonicalLink.getAttribute('href') || ''; } // Get h1 tags document.querySelectorAll('h1').forEach(h1 => { const text = h1.innerText.trim(); if (text) data.h1.push(text); }); // Get h2 tags document.querySelectorAll('h2').forEach(h2 => { const text = h2.innerText.trim(); if (text) data.h2.push(text); }); // Count images data.images = document.querySelectorAll('img').length; // Count links data.links = document.querySelectorAll('a').length; return data; }''') result = { "status": "success", "url": decoded_url, "seo": seo_data } return result except Exception as e: print(f"Error during SEO extraction: {e}") return {"status": "error", "url": decoded_url, "error": str(e)} # Perform the operation result = await safe_browser_operation(decoded_url, seo_operation) # Save to cache save_to_cache(decoded_url, "seo", result) return result except Exception as e: raise HTTPException(status_code=500, detail=str(e)) @app.get("/meta") async def extract_meta_tags(url: str, x_api_key: Optional[str] = Header(None)): """Extract meta tags from a website""" # Validate API key if not x_api_key or x_api_key != API_KEY: raise HTTPException(status_code=401, detail="Invalid API key") # Decode URL if it's encoded decoded_url = unquote(url) # Check cache first cached_result = get_cached_data(decoded_url, "meta") if cached_result: return cached_result try: print(f"Extracting meta tags from: {decoded_url}") # Define the operation to perform with the browser async def meta_operation(page): try: response = await page.goto(decoded_url, waitUntil='networkidle2', timeout=30000) # Extract all meta tags meta_tags = await page.evaluate('''() => { const metas = Array.from(document.querySelectorAll('meta')); return metas.map(meta => { const attributes = {}; Array.from(meta.attributes).forEach(attr => { attributes[attr.name] = attr.value; }); return attributes; }); }''') # Extract Open Graph tags og_tags = await page.evaluate('''() => { const ogTags = {}; document.querySelectorAll('meta[property^="og:"]').forEach(tag => { const property = tag.getAttribute('property'); ogTags[property] = tag.getAttribute('content'); }); return ogTags; }''') # Extract Twitter card tags twitter_tags = await page.evaluate('''() => { const twitterTags = {}; document.querySelectorAll('meta[name^="twitter:"]').forEach(tag => { const name = tag.getAttribute('name'); twitterTags[name] = tag.getAttribute('content'); }); return twitterTags; }''') result = { "status": "success", "url": decoded_url, "meta_tags": meta_tags, "open_graph": og_tags, "twitter_card": twitter_tags, "title": await page.title() } return result except Exception as e: print(f"Error during meta tag extraction: {e}") return {"status": "error", "url": decoded_url, "error": str(e)} # Perform the operation result = await safe_browser_operation(decoded_url, meta_operation) # Save to cache save_to_cache(decoded_url, "meta", result) return result except Exception as e: raise HTTPException(status_code=500, detail=str(e)) @app.get("/cache/clear") async def clear_cache(x_api_key: Optional[str] = Header(None)): """Clear the entire cache database""" # Validate API key if not x_api_key or x_api_key != API_KEY: raise HTTPException(status_code=401, detail="Invalid API key") conn = sqlite3.connect(DB_PATH) cursor = conn.cursor() cursor.execute("DELETE FROM cache") conn.commit() conn.close() return {"status": "success", "message": "Cache cleared successfully"} @app.get("/cache/stats") async def cache_stats(x_api_key: Optional[str] = Header(None)): """Get cache statistics""" # Validate API key if not x_api_key or x_api_key != API_KEY: raise HTTPException(status_code=401, detail="Invalid API key") try: conn = sqlite3.connect(DB_PATH) cursor = conn.cursor() # Get total count cursor.execute("SELECT COUNT(*) FROM cache") total_count = cursor.fetchone()[0] # Get count by route cursor.execute("SELECT route, COUNT(*) FROM cache GROUP BY route") route_counts = dict(cursor.fetchall()) # Get oldest and newest entries cursor.execute("SELECT MIN(timestamp), MAX(timestamp) FROM cache") min_time, max_time = cursor.fetchone() conn.close() return { "total_entries": total_count, "route_counts": route_counts, "oldest_entry": min_time, "newest_entry": max_time } except Exception as e: raise HTTPException(status_code=500, detail=str(e)) @asynccontextmanager async def get_browser(): """Get a browser from the pool or create a new one if needed""" browser = None try: # Try to get a browser from the pool try: browser = await browser_pool.get_nowait() except asyncio.QueueEmpty: # If pool is empty, create a new browser if under the limit async with browser_lock: if browser_pool.qsize() + 1 <= MAX_BROWSERS: browser = await create_browser() else: # If at limit, wait for a browser to become available browser = await browser_pool.get() yield browser finally: # Return browser to pool if it's still viable if browser: try: # Quick check if browser is still usable await browser.pages() # Check if browser is too old if time.time() - browser_creation_times.get(browser, 0) > BROWSER_TTL: await browser.close() if browser in browser_creation_times: del browser_creation_times[browser] browser = await create_browser() await browser_pool.put(browser) except Exception: # If browser is not usable, close it and create a new one try: await browser.close() except: pass if browser in browser_creation_times: del browser_creation_times[browser] browser = await create_browser() await browser_pool.put(browser) if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8000)