150 lines
4.3 KiB
Python
150 lines
4.3 KiB
Python
from fastapi import FastAPI, HTTPException, Header
|
|
from pyppeteer import launch
|
|
import os
|
|
import asyncio
|
|
from typing import Optional
|
|
from urllib.parse import unquote
|
|
import logging
|
|
import gc
|
|
|
|
# Configure logging
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
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")
|
|
|
|
# Browser pool
|
|
browser = None
|
|
|
|
# Browser configuration
|
|
BROWSER_ARGS = [
|
|
'--no-sandbox',
|
|
'--disable-setuid-sandbox',
|
|
'--disable-dev-shm-usage', # Overcome limited resource problems
|
|
'--disable-gpu', # Disable GPU hardware acceleration
|
|
'--disable-extensions', # Disable extensions
|
|
'--disable-audio-output', # Disable audio
|
|
'--disable-background-networking',
|
|
'--disable-background-timer-throttling',
|
|
'--disable-backgrounding-occluded-windows',
|
|
'--disable-breakpad',
|
|
'--disable-component-extensions-with-background-pages',
|
|
'--disable-features=TranslateUI,BlinkGenPropertyTrees',
|
|
'--disable-ipc-flooding-protection',
|
|
'--disable-renderer-backgrounding',
|
|
'--enable-features=NetworkService,NetworkServiceInProcess',
|
|
'--mute-audio',
|
|
'--memory-pressure-off',
|
|
]
|
|
|
|
# Browser launch options
|
|
BROWSER_OPTIONS = {
|
|
'headless': True,
|
|
'executablePath': '/usr/bin/google-chrome',
|
|
'args': BROWSER_ARGS,
|
|
'handleSIGINT': False,
|
|
'handleSIGTERM': False,
|
|
'handleSIGHUP': False,
|
|
}
|
|
|
|
async def get_browser():
|
|
"""Get or create a browser instance"""
|
|
global browser
|
|
if browser is None or not browser.isConnected():
|
|
logger.info("Launching new browser instance")
|
|
browser = await launch(**BROWSER_OPTIONS)
|
|
return browser
|
|
|
|
async def wait_for_network_idle(page):
|
|
"""Wait until no network requests are in flight"""
|
|
try:
|
|
await page.waitForNetworkIdle(idleTime=500, timeout=30000)
|
|
except Exception as e:
|
|
logger.warning(f"Network idle timeout: {str(e)}")
|
|
|
|
@app.on_event("startup")
|
|
async def startup_event():
|
|
"""Initialize browser on startup"""
|
|
await get_browser()
|
|
|
|
@app.on_event("shutdown")
|
|
async def shutdown_event():
|
|
"""Close browser on shutdown"""
|
|
global browser
|
|
if browser:
|
|
logger.info("Closing browser on shutdown")
|
|
await browser.close()
|
|
browser = None
|
|
|
|
@app.head("/")
|
|
async def health_check():
|
|
return {"status": "ok"}
|
|
|
|
@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")
|
|
|
|
page = None
|
|
try:
|
|
# Decode URL if it's encoded
|
|
decoded_url = unquote(url)
|
|
logger.info(f"Visiting URL: {decoded_url}")
|
|
|
|
# Get browser instance
|
|
browser_instance = await get_browser()
|
|
|
|
# Create new page with timeout
|
|
page = await browser_instance.newPage()
|
|
|
|
# Set page timeout and resource limits
|
|
await page.setDefaultNavigationTimeout(60000) # 60 seconds timeout
|
|
await page.setRequestInterception(True)
|
|
|
|
# Block unnecessary resources
|
|
async def intercept_request(request):
|
|
if request.resourceType in ['image', 'media', 'font']:
|
|
await request.abort()
|
|
else:
|
|
await request.continue_()
|
|
|
|
page.on('request', intercept_request)
|
|
|
|
# Navigate to URL and wait for network idle
|
|
await page.goto(decoded_url, waitUntil='networkidle2', timeout=60000)
|
|
|
|
# Get page content
|
|
content = await page.content()
|
|
|
|
# Close page to free up resources
|
|
await page.close()
|
|
|
|
# Force garbage collection
|
|
gc.collect()
|
|
|
|
return {"status": "success", "content": content}
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error processing request: {str(e)}")
|
|
# Ensure browser is still running
|
|
if browser is None or not browser.isConnected():
|
|
await get_browser()
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
finally:
|
|
# Always ensure page is closed
|
|
if page:
|
|
try:
|
|
await page.close()
|
|
except:
|
|
pass
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
uvicorn.run(app, host="0.0.0.0", port=8000)
|