Files
projects/Dockers/puppeteer-api/main.py
T
Bram c3fd4b3121
Build and Push Docker Images / build-and-push (push) Successful in 18s
nope
2025-04-30 21:01:06 +02:00

454 lines
14 KiB
Python

from fastapi import FastAPI, HTTPException, Header
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
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")
# 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'
# 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 36 hours
def get_cached_data(url, route):
conn = sqlite3.connect(DB_PATH)
cursor = conn.cursor()
cache_expiry = int(time.time()) - (36 * 60 * 60) # 36 hours in 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}")
# Initialize database on startup
init_db()
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"}
@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
browser = None
try:
print(decoded_url)
# Launch browser using installed Chrome
browser = await launch(
headless=True,
executablePath='/usr/bin/google-chrome',
args=['--no-sandbox', '--disable-setuid-sandbox']
)
# Create new page
page = await browser.newPage()
# Set custom user agent
await page.setUserAgent(CUSTOM_USER_AGENT)
# Navigate to URL and wait for network idle
await page.goto(decoded_url, waitUntil='networkidle2')
# Get page content
content = await page.content()
# Close page
await page.close()
result = {"status": "success", "content": content}
# Save to cache
save_to_cache(decoded_url, "visit", result)
return result
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
finally:
# Ensure browser is closed even if an error occurs
if browser:
await browser.close()
@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
browser = None
try:
print(f"Extracting SEO from: {decoded_url}")
# Launch browser using installed Chrome
browser = await launch(
headless=True,
executablePath='/usr/bin/google-chrome',
args=['--no-sandbox', '--disable-setuid-sandbox']
)
# Create new page
page = await browser.newPage()
# Set custom user agent
await page.setUserAgent(CUSTOM_USER_AGENT)
# Navigate to URL and wait for network idle
await page.goto(decoded_url, waitUntil='networkidle2')
# Extract JSON-LD data
json_ld_data = await page.evaluate('''() => {
const jsonLdScripts = Array.from(document.querySelectorAll('script[type="application/ld+json"]'));
return jsonLdScripts.map(script => {
try {
return JSON.parse(script.textContent);
} catch (e) {
return { error: "Failed to parse JSON-LD", content: script.textContent };
}
});
}''')
# Extract meta tags
meta_tags = await page.evaluate('''() => {
const metaTags = Array.from(document.querySelectorAll('meta'));
return metaTags.map(tag => {
const attributes = Array.from(tag.attributes);
const result = {};
attributes.forEach(attr => {
result[attr.name] = attr.value;
});
return result;
});
}''')
# Extract title
title = await page.evaluate('document.title')
# Extract canonical URL
canonical_url = await page.evaluate('''() => {
const link = document.querySelector('link[rel="canonical"]');
return link ? link.href : null;
}''')
# Extract Open Graph data
og_data = await page.evaluate('''() => {
const ogTags = Array.from(document.querySelectorAll('meta[property^="og:"]'));
const result = {};
ogTags.forEach(tag => {
const property = tag.getAttribute('property');
const content = tag.getAttribute('content');
result[property] = content;
});
return result;
}''')
# Extract Twitter Card data
twitter_data = await page.evaluate('''() => {
const twitterTags = Array.from(document.querySelectorAll('meta[name^="twitter:"]'));
const result = {};
twitterTags.forEach(tag => {
const name = tag.getAttribute('name');
const content = tag.getAttribute('content');
result[name] = content;
});
return result;
}''')
# Close page
await page.close()
# Compile all SEO data
seo_data = {
"title": title,
"canonical_url": canonical_url,
"json_ld": json_ld_data,
"meta_tags": meta_tags,
"open_graph": og_data,
"twitter_card": twitter_data
}
result = {"status": "success", "seo_data": seo_data}
# Save to cache
save_to_cache(decoded_url, "seo", result)
return result
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
finally:
# Ensure browser is closed even if an error occurs
if browser:
await browser.close()
@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
browser = None
try:
print(f"Extracting meta tags from: {decoded_url}")
# Launch browser using installed Chrome
browser = await launch(
headless=True,
executablePath='/usr/bin/google-chrome',
args=['--no-sandbox', '--disable-setuid-sandbox']
)
# Create new page
page = await browser.newPage()
# Set custom user agent
await page.setUserAgent(CUSTOM_USER_AGENT)
# Navigate to URL and wait for network idle
await page.goto(decoded_url, waitUntil='networkidle2')
# Extract meta tags
meta_tags = await page.evaluate('''() => {
const metaTags = Array.from(document.querySelectorAll('meta'));
return metaTags.map(tag => {
const attributes = Array.from(tag.attributes);
const result = {};
attributes.forEach(attr => {
result[attr.name] = attr.value;
});
return result;
});
}''')
# Extract title
title = await page.evaluate('document.title')
# Close page
await page.close()
result = {
"status": "success",
"url": decoded_url,
"title": title,
"meta_tags": meta_tags
}
# Save to cache
save_to_cache(decoded_url, "meta", result)
return result
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
finally:
# Ensure browser is closed even if an error occurs
if browser:
await browser.close()
@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")
conn = sqlite3.connect(DB_PATH)
cursor = conn.cursor()
# Get total entries
cursor.execute("SELECT COUNT(*) FROM cache")
total_entries = cursor.fetchone()[0]
# Get entries by route
cursor.execute("SELECT route, COUNT(*) FROM cache GROUP BY route")
routes = {route: count for route, count in cursor.fetchall()}
# Get recent entries (last 24 hours)
recent_timestamp = int(time.time()) - (24 * 60 * 60)
cursor.execute("SELECT COUNT(*) FROM cache WHERE timestamp > ?", (recent_timestamp,))
recent_entries = cursor.fetchone()[0]
conn.close()
return {
"status": "success",
"stats": {
"total_entries": total_entries,
"entries_by_route": routes,
"recent_entries": recent_entries
}
}
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=8000)